Release GenAI Conventions + Gemini & VertexAI Instrumentation

README.md

@@ -240,12 +240,14 @@ Langtrace automatically captures traces from the following vendors:
 | Cohere       | LLM             | :white_check_mark: | :white_check_mark:              |
 | Groq         | LLM             | :x:                | :white_check_mark:              |
 | Perplexity   | LLM             | :white_check_mark: | :white_check_mark:              |
+| Gemini       | LLM             | :x:                | :white_check_mark:              |
 | Langchain    | Framework       | :x:                | :white_check_mark:              |
 | LlamaIndex   | Framework       | :white_check_mark: | :white_check_mark:              |
 | Langgraph    | Framework       | :x:                | :white_check_mark:              |
 | DSPy         | Framework       | :x:                | :white_check_mark:              |
 | CrewAI       | Framework       | :x:                | :white_check_mark:              |
 | Ollama       | Framework       | :x:                | :white_check_mark:              |
+| VertexAI     | Framework       | :x:                | :white_check_mark:              |
 | Pinecone     | Vector Database | :white_check_mark: | :white_check_mark:              |
 | ChromaDB     | Vector Database | :white_check_mark: | :white_check_mark:              |
 | QDrant       | Vector Database | :white_check_mark: | :white_check_mark:              |

pyproject.toml

@@ -18,7 +18,7 @@ classifiers=[
     "Operating System :: OS Independent",
 ]
 dependencies = [
-  'trace-attributes>=5.0.0,<6.0.0',
+  'trace-attributes>=6.0.0,<7.0.0',
   'opentelemetry-api>=1.25.0',
   'opentelemetry-sdk>=1.25.0',
   'opentelemetry-instrumentation>=0.46b0',
@@ -49,7 +49,10 @@ dev = [
     "cohere",
     "qdrant_client",
     "weaviate-client",
-    "ollama"
+    "ollama",
+    "groq",
+    "google-generativeai",
+    "google-cloud-aiplatform"
 ]
 
 test = [

src/examples/cohere_example/chat.py

@@ -23,6 +23,7 @@ def chat_comp():
                 "message": "The man who is widely credited with discovering gravity is Sir Isaac Newton",
             },
         ],
+        k=3,
         message="Tell me a story in 3 sentences or less?",
         preamble="answer like a pirate",
         # perform web search before answering the question. You can also use your own custom connector.

src/examples/cohere_example/chat_stream.py

@@ -18,6 +18,9 @@ def chat_stream():
         message="Tell me a short story in 2 lines",
         preamble="Respond like a pirate",
         max_tokens=100,
+        k=3,
+        p=0.9,
+        temperature=0.5,
     ):
         if event.event_type == "text-generation":
             result.append(event.text)

src/examples/gemini_example/__init__.py

@@ -0,0 +1,6 @@
+from .main import basic
+
+
+class GeminiRunner:
+    def run(self):
+        basic()

src/examples/gemini_example/function_tools.py

@@ -0,0 +1,62 @@
+tools = [
+    {
+        "function_declarations": [
+            {
+                "name": "find_movies",
+                "description": "find movie titles currently playing in theaters based on any description, genre, title words, etc.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                        },
+                        "description": {
+                            "type": "string",
+                            "description": "Any kind of description including category or genre, title words, attributes, etc.",
+                        },
+                    },
+                    "required": ["description"],
+                },
+            },
+            {
+                "name": "find_theaters",
+                "description": "find theaters based on location and optionally movie title which is currently playing in theaters",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                        },
+                        "movie": {"type": "string", "description": "Any movie title"},
+                    },
+                    "required": ["location"],
+                },
+            },
+            {
+                "name": "get_showtimes",
+                "description": "Find the start times for movies playing in a specific theater",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                        },
+                        "movie": {"type": "string", "description": "Any movie title"},
+                        "theater": {
+                            "type": "string",
+                            "description": "Name of the theater",
+                        },
+                        "date": {
+                            "type": "string",
+                            "description": "Date for requested showtime",
+                        },
+                    },
+                    "required": ["location", "movie", "theater", "date"],
+                },
+            },
+        ]
+    }
+]

src/examples/gemini_example/main.py

@@ -0,0 +1,91 @@
+from langtrace_python_sdk import langtrace
+import google.generativeai as genai
+from dotenv import load_dotenv
+import os
+import asyncio
+import pathlib
+from .function_tools import tools
+
+load_dotenv()
+
+langtrace.init(write_spans_to_console=False, batch=False)
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+
+
+async def async_demo():
+    task1 = asyncio.create_task(async_generate())
+    task2 = asyncio.create_task(async_generate(stream=True))
+    return await asyncio.gather(task1, task2)
+
+
+def basic():
+    generate()
+    generate(stream=True, with_tools=True)
+
+    # image_to_text()
+    # audio_to_text()
+    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."
+    )
+
+    response = model.generate_content(
+        "Write a story about a AI and magic",
+        stream=stream,
+        tools=tools if with_tools else None,
+    )
+    if stream:
+        for res in response:
+            if res.text:
+                print(res.text)
+    else:
+        print(response.text)
+
+
+async def async_generate(stream=False):
+    model = genai.GenerativeModel(
+        "gemini-1.5-pro", system_instruction="You are a cat. Your name is Neko."
+    )
+    response = await model.generate_content_async(
+        "Write a story about a AI and magic", stream=stream
+    )
+    if stream:
+        async for chunk in response:
+            if chunk.text:
+                print(chunk.text)
+    else:
+        print(response.text)
+
+
+def image_to_text(stream=False):
+    model = genai.GenerativeModel("gemini-1.5-flash")
+    image1 = {
+        "mime_type": "image/jpeg",
+        "data": pathlib.Path("src/examples/gemini_example/jetpack.jpg").read_bytes(),
+    }
+
+    prompt = "Describe me this picture. What do you see in it."
+    response = model.generate_content([prompt, image1], stream=stream)
+    if stream:
+        for res in response:
+            print(res.text)
+    else:
+        print(response.text)
+
+
+# def audio_to_text(stream=False):
+#     model = genai.GenerativeModel("gemini-1.5-flash")
+#     audio = genai.upload_file(
+#         pathlib.Path("src/examples/gemini_example/voice_note.mp3")
+#     )
+
+#     prompt = "Summarize this voice recording."
+#     response = model.generate_content([prompt, audio], stream=stream)
+#     if stream:
+#         for res in response:
+#             print(res.text)
+#     else:
+#         print(response.text)

src/examples/langchain_example/__init__.py

@@ -1,10 +1,18 @@
 from .basic import basic_app, rag, load_and_split
 from langtrace_python_sdk import with_langtrace_root_span
 
+from .groq_example import groq_basic, groq_streaming
+
 
 class LangChainRunner:
     @with_langtrace_root_span("LangChain")
     def run(self):
         basic_app()
         rag()
         load_and_split()
+
+
+class GroqRunner:
+    @with_langtrace_root_span("Groq")
+    def run(self):
+        groq_streaming()

src/examples/langchain_example/groq_example.py

@@ -1,6 +1,7 @@
 from dotenv import find_dotenv, load_dotenv
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_groq import ChatGroq
+from groq import Groq
 
 _ = load_dotenv(find_dotenv())
 
@@ -12,21 +13,33 @@
 
 langtrace.init()
 
+client = Groq()
 
-def groq_example():
 
-    chat = ChatGroq(temperature=0, model_name="mixtral-8x7b-32768")
-
-    system = "You are a helpful assistant."
-    human = "{text}"
-    prompt = ChatPromptTemplate.from_messages([("system", system), ("human", human)])
-
-    chain = prompt | chat
-    result = chain.invoke(
-        {"text": "Explain the importance of low latency LLMs in 2 sentences or less."}
+def groq_basic():
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "Explain the importance of low latency LLMs",
+            }
+        ],
+        stream=False,
+        model="llama3-8b-8192",
     )
-    # print(result)
-    return result
-
-
-groq_example()
+    return chat_completion
+
+
+def groq_streaming():
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "Explain the importance of low latency LLMs",
+            }
+        ],
+        stream=True,
+        model="llama3-8b-8192",
+    )
+    for chunk in chat_completion:
+        print(chunk)

src/examples/ollama_example/basic.py

@@ -17,6 +17,7 @@ def chat():
                 "content": "hi",
             },
         ],
+        options={"temperature": 0.5},
         stream=True,
     )
 

src/examples/openai_example/__init__.py

@@ -10,6 +10,7 @@ def run(self):
             run_conversation as run_conversation_streaming,
         )
         from .chat_completion import chat_completion as chat_completion_example
+
         from .embeddings_create import embeddings_create as embeddings_create_example
         from .function_calling import function_calling as function_example
         from .images_edit import image_edit

src/examples/openai_example/async_tool_calling_nonstreaming.py

@@ -8,7 +8,7 @@
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init()
+langtrace.init(write_spans_to_console=True)
 
 client = AsyncOpenAI()
 

src/examples/openai_example/chat_completion.py

@@ -9,7 +9,7 @@
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init(write_spans_to_console=False)
+langtrace.init(write_spans_to_console=True)
 client = OpenAI()
 
 

src/examples/openai_example/embeddings_create.py

@@ -16,5 +16,6 @@ def embeddings_create():
     result = client.embeddings.create(
         model="text-embedding-ada-002",
         input="Once upon a time, there was a pirate.",
+        encoding_format="float",
     )
     return result

src/examples/openai_example/images_edit.py

@@ -23,8 +23,8 @@ def image_edit():
 
     response = client.images.edit(
         model="dall-e-2",
-        image=open("./resources/lounge_flamingo.png", "rb"),
-        mask=open("./resources/mask.png", "rb"),
+        image=open("src/examples/openai_example/resources/lounge_flamingo.png", "rb"),
+        mask=open("src/examples/openai_example/resources/mask.png", "rb"),
         prompt="A sunlit indoor lounge area with a pool and duck standing in side with flamingo.",
         n=1,
         size="1024x1024",

src/examples/vertexai_example/__init__.py

@@ -0,0 +1,6 @@
+from .main import basic
+
+
+class VertexAIRunner:
+    def run(self):
+        basic()