feat: Ibm watsonx integration

examples/anthropic/async-streaming.py

@@ -0,0 +1,27 @@
+import os
+import asyncio
+from anthropic import AsyncAnthropic
+import lunary
+
+client = AsyncAnthropic(
+    api_key=os.environ.get("ANTHROPIC_API_KEY"),  # This is the default and can be omitted
+)
+lunary.monitor(client)
+
+
+async def main() -> None:
+    stream = await client.messages.create(
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Hello, Claude",
+            }
+        ],
+        model="claude-3-opus-20240229",
+    )
+    for event in stream:
+        pass
+
+
+asyncio.run(main())

examples/anthropic/async.py

@@ -0,0 +1,26 @@
+import os
+import asyncio
+from anthropic import AsyncAnthropic
+import lunary
+
+client = AsyncAnthropic(
+    api_key=os.environ.get("ANTHROPIC_API_KEY"),  # This is the default and can be omitted
+)
+lunary.monitor(client)
+
+
+async def main() -> None:
+    message = await client.messages.create(
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Hello, Claude",
+            }
+        ],
+        model="claude-3-opus-20240229",
+    )
+    print(message.content)
+
+
+asyncio.run(main())

examples/anthropic/basic.py

@@ -0,0 +1,22 @@
+import os
+from anthropic import Anthropic
+import lunary
+
+client = Anthropic(
+    api_key=os.environ.get("ANTHROPIC_API_KEY"),  
+)
+lunary.monitor(client)
+
+
+message = client.messages.create(
+    max_tokens=1024,
+    messages=[
+        {
+            "role": "user",
+            "content": "Hello, Claude",
+        }
+    ],
+    model="claude-3-opus-20240229",
+)
+
+print(message.ro

examples/anthropic/streaming.py

@@ -0,0 +1,26 @@
+import os
+from anthropic import Anthropic
+import lunary
+
+client = Anthropic(
+    api_key=os.environ.get("ANTHROPIC_API_KEY"),  
+)
+lunary.monitor(client)
+
+
+stream = client.messages.create(
+    max_tokens=1024,
+    messages=[
+        {
+            "role": "user",
+            "content": "Hello, Claude",
+        }
+    ],
+    model="claude-3-opus-20240229",
+    stream=True
+)
+
+for event in stream:
+  pass
+
+

examples/anthropic/tool-call.py

@@ -0,0 +1,86 @@
+from anthropic import Anthropic 
+import lunary
+
+client = Anthropic()
+lunary.monitor(client)
+
+response = client.messages.create(
+    model="claude-3-5-sonnet-20241022",
+    max_tokens=1024,
+    tools=[
+        {
+            "name": "get_weather",
+            "description": "Get the current weather in a given location",
+            "input_schema": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    }
+                },
+                "required": ["location"],
+            },
+        }
+    ],
+    messages=[{"role": "user", "content": "What's the weather like in San Francisco?"}],
+)
+
+response = client.messages.create(
+    model="claude-3-5-sonnet-20241022",
+    max_tokens=1024,
+    tools=[
+        {
+            "name": "get_weather",
+            "description": "Get the current weather in a given location",
+            "input_schema": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA"
+                    },
+                    "unit": {
+                        "type": "string",
+                        "enum": ["celsius", "fahrenheit"],
+                        "description": "The unit of temperature, either 'celsius' or 'fahrenheit'"
+                    }
+                },
+                "required": ["location"]
+            }
+        }
+    ],
+    messages=[
+        {
+            "role": "user",
+            "content": "What's the weather like in San Francisco?"
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "text",
+                    "text": "<thinking>I need to use get_weather, and the user wants SF, which is likely San Francisco, CA.</thinking>"
+                },
+                {
+                    "type": "tool_use",
+                    "id": "toolu_01A09q90qw90lq917835lq9",
+                    "name": "get_weather",
+                    "input": {"location": "San Francisco, CA", "unit": "celsius"}
+                }
+            ]
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": "toolu_01A09q90qw90lq917835lq9", # from the API response
+                    "content": "65 degrees" # from running your tool
+                }
+            ]
+        }
+    ]
+)
+
+print(response)

examples/ibm/basic.py

@@ -0,0 +1,24 @@
+from ibm_watsonx_ai import Credentials
+from ibm_watsonx_ai.foundation_models import ModelInference
+import lunary
+import json
+
+model = ModelInference(
+    model_id="meta-llama/llama-3-1-8b-instruct",
+    credentials=Credentials(
+        api_key = "pfc6XWOR0oX_oZY-c4axG-JKKzMhKz1NF04g9C8Idfz2",
+        url = "https://us-south.ml.cloud.ibm.com"),
+        project_id="c36b3e63-c64c-4f2c-8885-933244642424"
+    )
+
+lunary.monitor_ibm(model)
+
+print(model.model_id)
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Who won the world series in 2020?"}
+]
+response = model.chat(messages=messages)
+
+# print(json.dumps(response, indent=4))
+# print(generated_response['choices'][0]['message']['content'])

examples/tools.py

@@ -0,0 +1,28 @@
+from openai import OpenAI
+import lunary
+
+client = OpenAI()
+lunary.monitor(client)
+
+tools = [
+  {
+      "type": "function",
+      "function": {
+          "name": "get_weather",
+          "parameters": {
+              "type": "object",
+              "properties": {
+                  "location": {"type": "string"}
+              },
+          },
+      },
+  }
+]
+
+completion = client.chat.completions.create(
+  model="gpt-4o",
+  messages=[{"role": "user", "content": "What's the weather like in Paris today?"}],
+  tools=tools,
+)
+
+print(completion.choices[0].message.tool_calls)

lunary/__init__.py

@@ -15,6 +15,7 @@
 from .exceptions import *
 from .parsers import default_input_parser, default_output_parser, filter_params, method_input_parser, PydanticHandler
 from .openai_utils import OpenAIUtils
+from .ibm_utils import IBMUtils
 from .event_queue import EventQueue
 from .thread import Thread
 from .utils import clean_nones, create_uuid_from_string
@@ -564,6 +565,15 @@ def async_stream_wrapper(*args, **kwargs):
     return wrapper
 
 
+def monitor_ibm(object):
+    object.chat = wrap(
+        object.chat,
+        "llm",
+        input_parser=IBMUtils.parse_input,
+        output_parser=IBMUtils.parse_output,
+        name=object.model_id,
+    )
+
 def monitor(object):
     try:
         openai_version = parse_version(version("openai"))

lunary/anthropic_utils.py

@@ -0,0 +1,83 @@
+import json, logging
+import anthropic
+
+logger = logging.getLogger(__name__)
+
+# TODO: make sure it's the correct list
+KWARGS_TO_CAPTURE = [
+    "frequency_penalty",
+    "functions",
+    "logit_bias",
+    "max_tokens",
+    "max_completion_tokens",
+    "n",
+    "presence_penalty",
+    "response_format",
+    "seed",
+    "stop",
+    "stream",
+    "audio",
+    "modalities",
+    "temperature",
+    "tool_choice",
+    "tools",
+    "tool_calls",
+    "top_p",
+    "top_k",
+    "top_logprobs",
+    "logprobs",
+    "prediction",
+    "service_tier",
+    "parallel_tool_calls"
+]
+
+class AnthropicUtils:
+    @staticmethod
+    def parse_message(message):
+        tool_calls = getattr(message, "tool_calls")
+
+        if tool_calls is not None:
+            tool_calls = [
+                json.loads(tool_calls.model_dump_json(index=2, exclude_unset=True))
+                for tool_calls in tool_calls 
+            ]
+        # TODO: audio?
+        # audio = getattr(message, "audio") 
+        # if audio is not None:
+        #     audio = json.loads(audio.model_dump_json(indent=2, exclude_unset=True))        
+
+        parsed_message = {
+            "role": getattr(message, "role"),
+            "content": getattr(message, "content"),
+            "refusal": getattr(message, "refusal"),
+            # TODO: "audio": audio?
+            # TODO: function_calls?
+            "tool_calls": getattr(message, "tool_calls")
+        }
+        return parsed_message
+
+
+    @staticmethod
+    def parse_input(*args, **kwargs):
+        try:
+            messages = [AnthropicUtils.parse_message(message) for message in kwargs["messages"]]
+            name = kwargs.get("model")
+            extra = {key: kwargs[key] for key in KWARGS_TO_CAPTURE if key in kwargs}
+
+            extra = {k: v for k, v in kwargs.items() if k in KWARGS_TO_CAPTURE}
+            return {"name": name, "input": messages, "extra": extra}
+        except Exception as e:
+            logger.error("Error parsing input: ", e)
+
+    @staticmethod
+    def parse_output(message, stream=False):
+        try:
+            parsed_output = {
+                "output": AnthropicUtils.parse_message(message),
+                "tokenUsage": {
+                    "prompt": getattr(message.usage, "input_tokens"),
+                    "completion": getattr(message.usage, "output_tokens") 
+                }
+            } 
+        except Exception as e:
+            logger.error("Error parsing output: ", e)