Merge branch 'development' of github.com:Scale3-Labs/langtrace-python-sdk into development

Author: karthikscale3

src/examples/awsbedrock_examples/__init__.py

@@ -1,10 +1,20 @@
-from examples.awsbedrock_examples.converse import use_converse
+from examples.awsbedrock_examples.converse import (
+    use_invoke_model_anthropic,
+    init_bedrock_langchain,
+    use_invoke_model_titan,
+    use_invoke_model_llama,
+)
 from langtrace_python_sdk import langtrace, with_langtrace_root_span
 
-langtrace.init()
-
 
 class AWSBedrockRunner:
     @with_langtrace_root_span("AWS_Bedrock")
     def run(self):
-        use_converse()
+
+        # use_converse_stream()
+        # use_converse()
+        # use_invoke_model_anthropic(stream=True)
+        # use_invoke_model_cohere()
+        # use_invoke_model_llama(stream=False)
+        # use_invoke_model_titan(stream=False)
+        init_bedrock_langchain()

src/examples/awsbedrock_examples/converse.py

@@ -1,34 +1,199 @@
-import os
 import boto3
+import json
 from langtrace_python_sdk import langtrace
+from dotenv import load_dotenv
+import botocore
+from langchain_aws import ChatBedrock
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False)
+
+brt = boto3.client("bedrock-runtime", region_name="us-east-1")
+brc = boto3.client("bedrock", region_name="us-east-1")
+
+
+def use_converse_stream():
+    model_id = "anthropic.claude-3-haiku-20240307-v1:0"
+    conversation = [
+        {
+            "role": "user",
+            "content": [{"text": "what is the capital of France?"}],
+        }
+    ]
+
+    try:
+        response = brt.converse_stream(
+            modelId=model_id,
+            messages=conversation,
+            inferenceConfig={"maxTokens": 4096, "temperature": 0},
+            additionalModelRequestFields={"top_k": 250},
+        )
+        # response_text = response["output"]["message"]["content"][0]["text"]
+        print(response)
+
+    except Exception as e:
+        print(f"ERROR: Can't invoke '{model_id}'. Reason: {e}")
+        exit(1)
 
-langtrace.init(api_key=os.environ["LANGTRACE_API_KEY"])
 
 def use_converse():
     model_id = "anthropic.claude-3-haiku-20240307-v1:0"
-    client = boto3.client(
-        "bedrock-runtime",
-        region_name="us-east-1",
-        aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
-        aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
-    )
     conversation = [
         {
             "role": "user",
-            "content": [{"text": "Write a story about a magic backpack."}],
+            "content": [{"text": "what is the capital of France?"}],
         }
     ]
 
     try:
-        response = client.converse(
+        response = brt.converse(
             modelId=model_id,
             messages=conversation,
-            inferenceConfig={"maxTokens":4096,"temperature":0},
-            additionalModelRequestFields={"top_k":250}
+            inferenceConfig={"maxTokens": 4096, "temperature": 0},
+            additionalModelRequestFields={"top_k": 250},
         )
         response_text = response["output"]["message"]["content"][0]["text"]
         print(response_text)
 
-    except (Exception) as e:
+    except Exception as e:
         print(f"ERROR: Can't invoke '{model_id}'. Reason: {e}")
-        exit(1)
+        exit(1)
+
+
+def get_foundation_models():
+    for model in brc.list_foundation_models()["modelSummaries"]:
+        print(model["modelId"])
+
+
+# Invoke Model API
+# Amazon Titan Models
+def use_invoke_model_titan(stream=False):
+    try:
+        prompt_data = "what's the capital of France?"
+        body = json.dumps(
+            {
+                "inputText": prompt_data,
+                "textGenerationConfig": {
+                    "maxTokenCount": 1024,
+                    "topP": 0.95,
+                    "temperature": 0.2,
+                },
+            }
+        )
+        modelId = "amazon.titan-text-express-v1"  # "amazon.titan-tg1-large"
+        accept = "application/json"
+        contentType = "application/json"
+
+        if stream:
+
+            response = brt.invoke_model_with_response_stream(
+                body=body, modelId=modelId, accept=accept, contentType=contentType
+            )
+            # Extract and print the response text in real-time.
+            for event in response["body"]:
+                chunk = json.loads(event["chunk"]["bytes"])
+                if "outputText" in chunk:
+                    print(chunk["outputText"], end="")
+
+        else:
+            response = brt.invoke_model(
+                body=body, modelId=modelId, accept=accept, contentType=contentType
+            )
+            response_body = json.loads(response.get("body").read())
+
+    except botocore.exceptions.ClientError as error:
+
+        if error.response["Error"]["Code"] == "AccessDeniedException":
+            print(
+                f"\x1b[41m{error.response['Error']['Message']}\
+                    \nTo troubeshoot this issue please refer to the following resources.\
+                    \nhttps://docs.aws.amazon.com/IAM/latest/UserGuide/troubleshoot_access-denied.html\
+                    \nhttps://docs.aws.amazon.com/bedrock/latest/userguide/security-iam.html\x1b[0m\n"
+            )
+
+        else:
+            raise error
+
+
+# Anthropic Models
+def use_invoke_model_anthropic(stream=False):
+    body = json.dumps(
+        {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 1024,
+            "temperature": 0.1,
+            "top_p": 0.9,
+            "messages": [{"role": "user", "content": "Hello, Claude"}],
+        }
+    )
+    modelId = "anthropic.claude-v2"
+    accept = "application/json"
+    contentType = "application/json"
+
+    if stream:
+        response = brt.invoke_model_with_response_stream(body=body, modelId=modelId)
+        stream_response = response.get("body")
+        if stream_response:
+            for event in stream_response:
+                chunk = event.get("chunk")
+                if chunk:
+                    # print(json.loads(chunk.get("bytes").decode()))
+                    pass
+
+    else:
+        response = brt.invoke_model(
+            body=body, modelId=modelId, accept=accept, contentType=contentType
+        )
+        response_body = json.loads(response.get("body").read())
+        # text
+        print(response_body.get("completion"))
+
+
+def use_invoke_model_llama(stream=False):
+    model_id = "meta.llama3-8b-instruct-v1:0"
+    prompt = "What is the capital of France?"
+    max_gen_len = 128
+    temperature = 0.1
+    top_p = 0.9
+
+    # Create request body.
+    body = json.dumps(
+        {
+            "prompt": prompt,
+            "max_gen_len": max_gen_len,
+            "temperature": temperature,
+            "top_p": top_p,
+        }
+    )
+
+    if stream:
+        response = brt.invoke_model_with_response_stream(body=body, modelId=model_id)
+        for event in response["body"]:
+            chunk = json.loads(event["chunk"]["bytes"])
+            if "generation" in chunk:
+                # print(chunk["generation"], end="")
+                pass
+    else:
+        response = brt.invoke_model(body=body, modelId=model_id)
+        response_body = json.loads(response.get("body").read())
+        return response_body
+
+
+# print(get_foundation_models())
+def use_invoke_model_cohere():
+    model_id = "cohere.command-r-plus-v1"
+    prompt = "What is the capital of France?"
+    body = json.dumps({"prompt": prompt, "max_tokens": 1024, "temperature": 0.1})
+    response = brt.invoke_model(body=body, modelId=model_id)
+    response_body = json.loads(response.get("body").read())
+    print(response_body)
+
+
+def init_bedrock_langchain(temperature=0.1):
+    chat = ChatBedrock(
+        model_id="anthropic.claude-v2",
+        streaming=True,
+        model_kwargs={"temperature": temperature},
+        region_name="us-east-1",
+    )
+    return chat.invoke("What is the capital of France?")

src/examples/gemini_example/main.py

@@ -20,16 +20,16 @@ async def async_demo():
 
 def basic():
     generate()
-    generate(stream=True, with_tools=True)
+    # generate(stream=True, with_tools=True)
 
     # image_to_text()
     # audio_to_text()
-    asyncio.run(async_demo())
+    # asyncio.run(async_demo())
 
 
 def generate(stream=False, with_tools=False):
     model = genai.GenerativeModel(
-        "gemini-1.5-pro", system_instruction="You are a cat. Your name is Neko."
+        "gemini-2.0-flash-exp", system_instruction="You are a cat. Your name is Neko."
     )
 
     response = model.generate_content(

src/examples/google_genai_example/__init__.py

@@ -0,0 +1,7 @@
+from .main import generate_content, generate_content_streaming
+
+
+class GoogleGenaiRunner:
+    def run(self):
+        # generate_content()
+        generate_content_streaming()

src/examples/google_genai_example/main.py

@@ -0,0 +1,27 @@
+from google import genai
+from dotenv import load_dotenv
+import os
+from langtrace_python_sdk import langtrace
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False)
+
+
+def generate_content():
+    # Only run this block for Google AI API
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-exp", contents="What is your name?"
+    )
+
+    print(response.text)
+
+
+def generate_content_streaming():
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content_stream(
+        model="gemini-2.0-flash-exp", contents="What is your name?"
+    )
+
+    for chunk in response:
+        pass

src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py

@@ -1,6 +1,10 @@
 from langtrace.trace_attributes import AWSBedrockMethods
 
 APIS = {
+    "INVOKE_MODEL": {
+        "METHOD": "aws_bedrock.invoke_model",
+        "ENDPOINT": "/invoke-model",
+    },
     "CONVERSE": {
         "METHOD": AWSBedrockMethods.CONVERSE.value,
         "ENDPOINT": "/converse",

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -10,6 +10,7 @@
 
 SERVICE_PROVIDERS = {
     "ANTHROPIC": "Anthropic",
+    "ARCH": "Arch",
     "AZURE": "Azure",
     "CHROMA": "Chroma",
     "CREWAI": "CrewAI",

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -24,6 +24,7 @@
 from .pymongo import PyMongoInstrumentation
 from .cerebras import CerebrasInstrumentation
 from .milvus import MilvusInstrumentation
+from .google_genai import GoogleGenaiInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -52,4 +53,5 @@
     "AWSBedrockInstrumentation",
     "CerebrasInstrumentation",
     "MilvusInstrumentation",
+    "GoogleGenaiInstrumentation",
 ]