Merge pull request #396 from Scale3-Labs/ali/s3en-2883-support-cerebras-python-library

Ali/s3en 2883 support cerebras python library

Author: alizenhom

src/examples/cerebras_example/__init__.py

@@ -0,0 +1,11 @@
+class CerebrasRunner:
+    def run(self):
+        from .main import (
+            completion_example,
+            completion_with_tools_example,
+            openai_cerebras_example,
+        )
+
+        completion_with_tools_example()
+        completion_example()
+        openai_cerebras_example()

src/examples/cerebras_example/main.py

@@ -0,0 +1,195 @@
+from langtrace_python_sdk import langtrace
+from cerebras.cloud.sdk import Cerebras
+from dotenv import load_dotenv
+import re
+import json
+from openai import OpenAI
+import os
+
+load_dotenv()
+
+langtrace.init()
+openai_client = OpenAI(
+    base_url="https://api.cerebras.ai/v1",
+    api_key=os.getenv("CEREBRAS_API_KEY"),
+)
+client = Cerebras()
+
+
+def openai_cerebras_example(stream=False):
+    completion = openai_client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "Why is fast inference important?",
+            }
+        ],
+        model="llama3.1-8b",
+        stream=stream,
+    )
+
+    if stream:
+        for chunk in completion:
+            print(chunk)
+    else:
+        return completion
+
+
+def completion_example(stream=False):
+    completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "Why is fast inference important?",
+            }
+        ],
+        model="llama3.1-8b",
+        stream=stream,
+    )
+
+    if stream:
+        for chunk in completion:
+            print(chunk)
+    else:
+        return completion
+
+
+def completion_with_tools_example(stream=False):
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant with access to a calculator. Use the calculator tool to compute mathematical expressions when needed.",
+        },
+        {"role": "user", "content": "What's the result of 15 multiplied by 7?"},
+    ]
+
+    response = client.chat.completions.create(
+        model="llama3.1-8b",
+        messages=messages,
+        tools=tools,
+        stream=stream,
+    )
+
+    if stream:
+        # Handle streaming response
+        full_content = ""
+        for chunk in response:
+            if chunk.choices[0].delta.tool_calls:
+                tool_call = chunk.choices[0].delta.tool_calls[0]
+                if hasattr(tool_call, "function"):
+                    if tool_call.function.name == "calculate":
+                        arguments = json.loads(tool_call.function.arguments)
+                        result = calculate(arguments["expression"])
+                        print(f"Calculation result: {result}")
+
+                        # Get final response with calculation result
+                        messages.append(
+                            {
+                                "role": "assistant",
+                                "content": None,
+                                "tool_calls": [
+                                    {
+                                        "function": {
+                                            "name": "calculate",
+                                            "arguments": tool_call.function.arguments,
+                                        },
+                                        "id": tool_call.id,
+                                        "type": "function",
+                                    }
+                                ],
+                            }
+                        )
+                        messages.append(
+                            {
+                                "role": "tool",
+                                "content": str(result),
+                                "tool_call_id": tool_call.id,
+                            }
+                        )
+
+                        final_response = client.chat.completions.create(
+                            model="llama3.1-70b", messages=messages, stream=True
+                        )
+
+                        for final_chunk in final_response:
+                            if final_chunk.choices[0].delta.content:
+                                print(final_chunk.choices[0].delta.content, end="")
+            elif chunk.choices[0].delta.content:
+                print(chunk.choices[0].delta.content, end="")
+                full_content += chunk.choices[0].delta.content
+    else:
+        # Handle non-streaming response
+        choice = response.choices[0].message
+        if choice.tool_calls:
+            function_call = choice.tool_calls[0].function
+            if function_call.name == "calculate":
+                arguments = json.loads(function_call.arguments)
+                result = calculate(arguments["expression"])
+                print(f"Calculation result: {result}")
+
+                messages.append(
+                    {
+                        "role": "assistant",
+                        "content": None,
+                        "tool_calls": [
+                            {
+                                "function": {
+                                    "name": "calculate",
+                                    "arguments": function_call.arguments,
+                                },
+                                "id": choice.tool_calls[0].id,
+                                "type": "function",
+                            }
+                        ],
+                    }
+                )
+                messages.append(
+                    {
+                        "role": "tool",
+                        "content": str(result),
+                        "tool_call_id": choice.tool_calls[0].id,
+                    }
+                )
+
+                final_response = client.chat.completions.create(
+                    model="llama3.1-70b",
+                    messages=messages,
+                )
+
+                if final_response:
+                    print(final_response.choices[0].message.content)
+                else:
+                    print("No final response received")
+        else:
+            print("Unexpected response from the model")
+
+
+def calculate(expression):
+    expression = re.sub(r"[^0-9+\-*/().]", "", expression)
+
+    try:
+        result = eval(expression)
+        return str(result)
+    except (SyntaxError, ZeroDivisionError, NameError, TypeError, OverflowError):
+        return "Error: Invalid expression"
+
+
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "calculate",
+            "description": "A calculator tool that can perform basic arithmetic operations. Use this when you need to compute mathematical expressions or solve numerical problems.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "expression": {
+                        "type": "string",
+                        "description": "The mathematical expression to evaluate",
+                    }
+                },
+                "required": ["expression"],
+            },
+        },
+    }
+]

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -35,6 +35,7 @@
     "AUTOGEN": "Autogen",
     "XAI": "XAI",
     "AWS_BEDROCK": "AWS Bedrock",
+    "CEREBRAS": "Cerebras",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -21,6 +21,7 @@
 from .aws_bedrock import AWSBedrockInstrumentation
 from .embedchain import EmbedchainInstrumentation
 from .litellm import LiteLLMInstrumentation
+from .cerebras import CerebrasInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -46,4 +47,5 @@
     "GeminiInstrumentation",
     "MistralInstrumentation",
     "AWSBedrockInstrumentation",
+    "CerebrasInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/cerebras/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import CerebrasInstrumentation
+
+__all__ = ["CerebrasInstrumentation"]

src/langtrace_python_sdk/instrumentation/cerebras/instrumentation.py

@@ -0,0 +1,54 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Collection
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from opentelemetry.semconv.schemas import Schemas
+from wrapt import wrap_function_wrapper
+from importlib_metadata import version as v
+from .patch import chat_completions_create, async_chat_completions_create
+
+
+class CerebrasInstrumentation(BaseInstrumentor):
+    """
+    The CerebrasInstrumentation class represents the Cerebras instrumentation
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["cerebras-cloud-sdk >= 1.0.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(
+            __name__, "", tracer_provider, schema_url=Schemas.V1_27_0.value
+        )
+        version = v("cerebras-cloud-sdk")
+
+        wrap_function_wrapper(
+            module="cerebras.cloud.sdk",
+            name="resources.chat.completions.CompletionsResource.create",
+            wrapper=chat_completions_create(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            module="cerebras.cloud.sdk",
+            name="resources.chat.completions.AsyncCompletionsResource.create",
+            wrapper=async_chat_completions_create(version, tracer),
+        )
+
+    def _uninstrument(self, **kwargs):
+        pass