Merge branch 'main' of github.com:Scale3-Labs/langtrace-python-sdk in…

…to release

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",

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -23,6 +23,7 @@
     "PPLX": "Perplexity",
     "QDRANT": "Qdrant",
     "WEAVIATE": "Weaviate",
+    "OLLAMA": "Ollama",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/constants/instrumentation/ollama.py

@@ -0,0 +1,7 @@
+APIS = {
+    "GENERATE": {
+        "METHOD": "generate",
+    },
+    "CHAT": {"METHOD": "chat"},
+    "EMBEDDINGS": {"METHOD": "embeddings"},
+}

src/langtrace_python_sdk/constants/instrumentation/pinecone.py

@@ -12,7 +12,7 @@
         "OPERATION": "query",
     },
     "DELETE": {
-        "METHOD": PineconeMethods.DELETE,
+        "METHOD": PineconeMethods.DELETE.value,
         "ENDPOINT": "/vectors/delete",
         "OPERATION": "delete",
     },

src/langtrace_python_sdk/extensions/langtrace_filesystem.py

@@ -15,7 +15,7 @@
 
 class OpenMode(str):
     def __init_subclass__(cls, **kwargs):
-        allowed_values = set(OpenTextMode.__args__) | set(OpenBinaryMode.__args__)
+        allowed_values = Union[set(OpenTextMode.__args__), set(OpenBinaryMode.__args__)]
         super().__init_subclass__(**kwargs)
 
         def __new__(cls, value):
@@ -27,6 +27,8 @@ def __new__(cls, value):
 
 
 class LangTraceFile(io.BytesIO):
+    _host: str = os.environ.get("LANGTRACE_API_HOST", None) or LANGTRACE_REMOTE_URL
+
     def __init__(self, fs: "LangTraceFileSystem", path: str, mode: OpenMode):
         super().__init__()
         self.fs = fs
@@ -50,34 +52,37 @@ def upload_to_server(self, file_data: bytes) -> None:
             log = file_data.decode("utf-8")
             eval_log = json.loads(log)
             data = {
-                "runId": eval_log['eval']['run_id'],
-                "taskId": eval_log['eval']['task_id'],
+                "runId": eval_log["eval"]["run_id"],
+                "taskId": eval_log["eval"]["task_id"],
                 "log": log,
             }
             if self.path is not None:
                 dataset_id = self.path.split("/")[0]
-                print(Fore.GREEN + f"Sending results to Langtrace for dataset: {dataset_id}" + Fore.RESET)
+                print(
+                    Fore.GREEN
+                    + f"Sending results to Langtrace for dataset: {dataset_id}"
+                    + Fore.RESET
+                )
                 data["datasetId"] = dataset_id
             else:
                 print(Fore.GREEN + "Sending results to Langtrace" + Fore.RESET)
             response = requests.post(
-                url=f"{LANGTRACE_REMOTE_URL}/api/run",
+                url=f"{self._host}/api/run",
                 data=json.dumps(data),
                 headers={
                     "Content-Type": "application/json",
-                    "x-api-key": os.environ.get("LANGTRACE_API_KEY")
+                    "x-api-key": os.environ.get("LANGTRACE_API_KEY"),
                 },
                 timeout=20,
             )
             response.raise_for_status()
-            print(
-                Fore.GREEN + "Results sent to Langtrace successfully." + Fore.RESET
-            )
+            print(Fore.GREEN + "Results sent to Langtrace successfully." + Fore.RESET)
         except requests.exceptions.RequestException as error:
             print(Fore.RED + f"Error reporting results: {error}" + Fore.RESET)
 
 
 class LangTraceFileSystem(AbstractFileSystem):
+    _host: str = os.environ.get("LANGTRACE_API_HOST", None) or LANGTRACE_REMOTE_URL
     protocol = "langtracefs"
     sep = "/"
 
@@ -89,9 +94,9 @@ def __init__(self, *args, **kwargs):
     def open(
         self,
         path: str,
-        mode: OpenTextMode | OpenBinaryMode = "rb",
+        mode: Union[OpenTextMode, OpenBinaryMode] = "rb",
         **kwargs,
-    ) -> Iterator[LangTraceFile | io.BytesIO]:
+    ) -> Iterator[Union[LangTraceFile, io.BytesIO]]:
         if "r" in mode:
             dataset_id = path
             # Fetch file from API and return a BytesIO object
@@ -104,21 +109,33 @@ def open(
 
     def fetch_file_from_api(self, dataset_id: str) -> bytes:
         try:
-            print(Fore.GREEN + f"Fetching dataset with id: {dataset_id} from Langtrace" + Fore.RESET)
+            print(
+                Fore.GREEN
+                + f"Fetching dataset with id: {dataset_id} from Langtrace"
+                + Fore.RESET
+            )
             response = requests.get(
-                url=f"{LANGTRACE_REMOTE_URL}/api/dataset/download?id={dataset_id}",
+                url=f"{self._host}/api/dataset/download?id={dataset_id}",
                 headers={
                     "Content-Type": "application/json",
-                    "x-api-key": os.environ.get("LANGTRACE_API_KEY")
+                    "x-api-key": os.environ.get("LANGTRACE_API_KEY"),
                 },
                 timeout=20,
             )
-            print(Fore.GREEN + f"Successfully fetched dataset with id: {dataset_id} from Langtrace" + Fore.RESET)
+            print(
+                Fore.GREEN
+                + f"Successfully fetched dataset with id: {dataset_id} from Langtrace"
+                + Fore.RESET
+            )
             response.raise_for_status()
             file_data = response.content
             return file_data
         except requests.exceptions.RequestException as error:
-            print(Fore.RED + f"Error fetching dataset with id: {dataset_id} from Langtrace: {error}" + Fore.RESET)
+            print(
+                Fore.RED
+                + f"Error fetching dataset with id: {dataset_id} from Langtrace: {error}"
+                + Fore.RESET
+            )
             return b""
 
     def makedirs(self, path: str, exist_ok: bool = False) -> None:
@@ -167,7 +184,9 @@ def _walk(self, path: str):
         if path in self.dirs:
             dirs = [d for d in self.dirs if d.startswith(path + self.sep)]
             files = [f for f in self.files if f.startswith(path + self.sep)]
-            yield path, [d.split(self.sep)[-1] for d in dirs], [f.split(self.sep)[-1] for f in files]
+            yield path, [d.split(self.sep)[-1] for d in dirs], [
+                f.split(self.sep)[-1] for f in files
+            ]
             for d in dirs:
                 yield from self._walk(d)
 

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -11,6 +11,7 @@
 from .pinecone import PineconeInstrumentation
 from .qdrant import QdrantInstrumentation
 from .weaviate import WeaviateInstrumentation
+from .ollama import OllamaInstrumentor
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -26,4 +27,5 @@
     "PineconeInstrumentation",
     "QdrantInstrumentation",
     "WeaviateInstrumentation",
+    "OllamaInstrumentor",
 ]