Merge pull request #207 from Scale3-Labs/development

Release Ollama + disable vendor specific functions

Author: alizenhom

pyproject.toml

@@ -49,6 +49,7 @@ dev = [
     "cohere",
     "qdrant_client",
     "weaviate-client",
+    "ollama"
 ]
 
 test = [

src/examples/inspect_ai_example/basic_eval.py

@@ -5,8 +5,7 @@
 from inspect_ai.scorer import model_graded_qa
 from inspect_ai.solver import chain_of_thought, generate, self_critique
 
-from langtrace_python_sdk.extensions.langtrace_filesystem import \
-    LangTraceFileSystem
+from langtrace_python_sdk.extensions.langtrace_filesystem import LangTraceFileSystem
 
 # from langtrace_python_sdk import langtrace
 
@@ -20,9 +19,6 @@
 def security_guide():
     return Task(
         dataset=csv_dataset("langtracefs://clxc2mxu6000lpc7ntsvcjvp9"),
-        plan=[
-            chain_of_thought(),
-            self_critique()
-        ],
-        scorer=model_graded_qa()
+        plan=[chain_of_thought(), self_critique()],
+        scorer=model_graded_qa(),
     )

src/examples/langchain_example/basic.py

@@ -8,14 +8,15 @@
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 
 from langtrace_python_sdk import langtrace
-from langtrace_python_sdk.utils.with_root_span import (
-    with_langtrace_root_span,
-    with_additional_attributes,
-)
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from openai import OpenAI
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init()
+langtrace.init(
+    write_spans_to_console=False,
+    disable_tracing_for_functions={"langchain": ["RunnableSequence.invoke"]},
+)
 
 
 def api_call_1():
@@ -29,7 +30,8 @@ def api_call_1():
     output_parser = StrOutputParser()
     chain = prompt | llm | output_parser
     res = chain.invoke({"input": "how can langsmith help with testing?"})
-    print(res)
+    # print(res)
+    return res
 
 
 def api_call_2():
@@ -43,13 +45,26 @@ def api_call_2():
     output_parser = StrOutputParser()
     chain = prompt | llm | output_parser
     res = chain.invoke({"input": "how can langsmith help with testing?"})
-    print(res)
+    # print(res)
+    return res
 
 
 @with_langtrace_root_span()
 def basic_app():
     api_call_1()
-    api_call_2()
+    # api_call_2()
+    client = OpenAI()
+    response = client.chat.completions.create(
+        model="gpt-4",
+        messages=[
+            {"role": "system", "content": "Talk like a pirate"},
+            {"role": "user", "content": "Tell me a story in 3 sentences or less."},
+        ],
+        # stream=True,
+        stream=False,
+    )
+
+    return response
 
 
 @with_langtrace_root_span()

src/examples/llamaindex_example/agent.py

@@ -12,7 +12,7 @@
 
 nest_asyncio.apply()
 
-langtrace.init(write_spans_to_console=False)
+langtrace.init()
 
 
 def multiply(a: int, b: int) -> int:

src/examples/llamaindex_example/basic.py

@@ -7,7 +7,11 @@
 _ = load_dotenv(find_dotenv())
 
 
-langtrace.init(write_spans_to_console=False)
+langtrace.init(
+    disable_tracing_for_functions={
+        "open_ai": ["openai.chat.completions.create"],
+    }
+)
 
 
 @with_langtrace_root_span()

src/examples/ollama_example/__init__.py

@@ -0,0 +1,14 @@
+from .basic import chat, async_chat, async_generate, generate, embed, async_embed
+from langtrace_python_sdk import with_langtrace_root_span
+import asyncio
+
+
+class OllamaRunner:
+    @with_langtrace_root_span("OllamaRunner")
+    def run(self):
+        chat()
+        generate()
+        embed()
+        asyncio.run(async_chat())
+        asyncio.run(async_generate())
+        asyncio.run(async_embed())

src/examples/ollama_example/basic.py

@@ -0,0 +1,50 @@
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
+import ollama
+from ollama import AsyncClient
+from dotenv import load_dotenv
+
+load_dotenv()
+
+langtrace.init(write_spans_to_console=False)
+
+
+def chat():
+    response = ollama.chat(
+        model="llama3",
+        messages=[
+            {
+                "role": "user",
+                "content": "hi",
+            },
+        ],
+        stream=True,
+    )
+
+    return response
+
+
+async def async_chat():
+    message = {"role": "user", "content": "Why is the sky blue?"}
+    return await AsyncClient().chat(model="llama3", messages=[message])
+
+
+def generate():
+    return ollama.generate(model="llama3", prompt="Why is the sky blue?")
+
+
+def async_generate():
+    return AsyncClient().generate(model="llama3", prompt="Why is the sky blue?")
+
+
+def embed():
+    return ollama.embeddings(
+        model="llama3",
+        prompt="cat",
+    )
+
+
+async def async_embed():
+    return await AsyncClient().embeddings(
+        model="llama3",
+        prompt="cat",
+    )

src/examples/openai_example/chat_completion.py

@@ -40,5 +40,5 @@ def chat_completion():
             ]
             result.append(content[0] if len(content) > 0 else "")
 
-    print("".join(result))
+    # print("".join(result))
     return response

src/examples/pinecone_example/basic.py

@@ -13,6 +13,7 @@
     with_additional_attributes,
 )
 from langtrace_python_sdk.utils.with_root_span import SendUserFeedback
+from opentelemetry.sdk.trace.export import ConsoleSpanExporter
 
 _ = load_dotenv(find_dotenv())
 langtrace.init()

src/langtrace_python_sdk/__init__.py

@@ -15,12 +15,14 @@
 """
 
 from langtrace_python_sdk import langtrace
-from langtrace_python_sdk.extensions.langtrace_filesystem import \
-    LangTraceFileSystem
+from langtrace_python_sdk.extensions.langtrace_filesystem import LangTraceFileSystem
 from langtrace_python_sdk.utils.prompt_registry import get_prompt_from_registry
 from langtrace_python_sdk.utils.with_root_span import (
-    SendUserFeedback, inject_additional_attributes, with_additional_attributes,
-    with_langtrace_root_span)
+    SendUserFeedback,
+    inject_additional_attributes,
+    with_additional_attributes,
+    with_langtrace_root_span,
+)
 
 __all__ = [
     "langtrace",