CrewAI enhancements and embedchain support (#313)

karthikscale3 · alizenhom · web-flow · commit cbb26659466b · 2024-08-30T10:13:16.000-07:00
* Minor fixes

* example update

* remove stream_usage

* Add embed chain support

* update README

* fix

* update version and deps

---------

Co-authored-by: Ali Waleed &lt;ali@scale3labs.com&gt;
diff --git a/README.md b/README.md
@@ -255,6 +255,7 @@ Langtrace automatically captures traces from the following vendors:
 | Ollama       | Framework       | :x:                | :white_check_mark:              |
 | VertexAI     | Framework       | :x:                | :white_check_mark:              |
 | Vercel AI SDK| Framework       | :white_check_mark: | :x:                             |
+| EmbedChain   | 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:              |
diff --git a/pyproject.toml b/pyproject.toml
@@ -55,6 +55,7 @@ dev = [
   "google-generativeai",
   "google-cloud-aiplatform",
   "mistralai",
+  "embedchain",
 ]
 
 test = ["pytest", "pytest-vcr", "pytest-asyncio"]
diff --git a/src/examples/crewai_example/simple_agent/agents.py b/src/examples/crewai_example/simple_agent/agents.py
@@ -1,15 +1,13 @@
 from crewai import Agent
-from langchain_openai import ChatOpenAI
 from langchain_anthropic import ChatAnthropic
 from langchain_cohere import ChatCohere
 from langchain_ollama import ChatOllama
+from langchain_openai import ChatOpenAI
 
 
 class PoetryAgents:
     def __init__(self):
-        self.open_ai = ChatOpenAI(
-            model_name="gpt-4", temperature=0.7, stream_usage=True
-        )
+        self.open_ai = ChatOpenAI(model_name="gpt-4", temperature=0.7)
         self.anthropic = ChatAnthropic(
             model_name="claude-3-5-sonnet-20240620", temperature=0.7
         )
@@ -30,3 +28,28 @@ def create_poet_agent(self):
             verbose=True,
             llm=self.open_ai,
         )
+    
+    def poet_agent_2(self):
+        return Agent(
+            role="Renaissance Poet",
+            backstory="""
+                    I am a Renaissance Poet. I am well-versed in the art of poetry and have a deep appreciation for the beauty of language and expression.
+                """,
+            goal="""Create a poem that is inspired by the works of the Renaissance poets""",
+            allow_delegation=False,
+            verbose=True,
+            llm=self.open_ai,
+        )
+
+    def poet_agent_3(self):
+        return Agent(
+            role="William Shakespeare",
+            backstory="""
+                    I am william shakespeare. I am an Expert in poetry writing and creative expression.
+                    I have been writing poetry for over 10 years and have published several collections.
+                """,
+            goal="""Create a poem that is inspired by the works of William Shakespeare""",
+            allow_delegation=False,
+            verbose=True,
+            llm=self.open_ai,
+        )
diff --git a/src/examples/crewai_example/simple_agent/main.py b/src/examples/crewai_example/simple_agent/main.py
@@ -1,8 +1,9 @@
+from agents import PoetryAgents
 from crewai import Crew
-from .agents import PoetryAgents
-from .tasks import PoetryTasks
-from langtrace_python_sdk import langtrace
 from dotenv import load_dotenv
+from tasks import PoetryTasks
+
+from langtrace_python_sdk import langtrace
 
 load_dotenv()
 langtrace.init()
@@ -17,10 +18,14 @@ def run(self):
         tasks = PoetryTasks()
 
         poetry_agent = agents.create_poet_agent()
+        poetry_agent_2 = agents.poet_agent_2()
+        poetry_agent_3 = agents.poet_agent_3()
 
         create_poem = tasks.create_poem(poetry_agent, self.topic)
+        create_poem_2 = tasks.create_poem(poetry_agent_2, self.topic)
+        create_poem_3 = tasks.create_poem(poetry_agent_3, self.topic)
 
-        crew = Crew(agents=[poetry_agent], tasks=[create_poem], verbose=True)
+        crew = Crew(agents=[poetry_agent], tasks=[create_poem], verbose=True, memory=True)
         res = crew.kickoff()
         return res
 
diff --git a/src/examples/embedchain_example/simple.py b/src/examples/embedchain_example/simple.py
@@ -0,0 +1,15 @@
+from dotenv import load_dotenv
+from embedchain import App
+from langtrace_python_sdk import langtrace
+
+load_dotenv()
+langtrace.init()
+
+app = App()
+app.reset()
+app.add("https://www.forbes.com/profile/elon-musk")
+app.add("https://en.wikipedia.org/wiki/Elon_Musk")
+res = app.query("What is the net worth of Elon Musk today?")
+print(res)
+re2 = app.search("Elon Musk")
+print(re2)
diff --git a/src/langtrace_python_sdk/constants/instrumentation/common.py b/src/langtrace_python_sdk/constants/instrumentation/common.py
@@ -30,6 +30,7 @@
     "VERTEXAI": "VertexAI",
     "GEMINI": "Gemini",
     "MISTRAL": "Mistral",
+    "EMBEDCHAIN": "Embedchain",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"
diff --git a/src/langtrace_python_sdk/constants/instrumentation/embedchain.py b/src/langtrace_python_sdk/constants/instrumentation/embedchain.py
@@ -0,0 +1,14 @@
+APIS = {
+    "ADD": {
+        "METHOD": "embedchain.add",
+        "OPERATION": "add",
+    },
+    "QUERY": {
+        "METHOD": "embedchain.query",
+        "OPERATION": "query",
+    },
+    "SEARCH": {
+        "METHOD": "embedchain.search",
+        "OPERATION": "search",
+    },
+}
diff --git a/src/langtrace_python_sdk/instrumentation/__init__.py b/src/langtrace_python_sdk/instrumentation/__init__.py
@@ -17,12 +17,14 @@
 from .vertexai import VertexAIInstrumentation
 from .gemini import GeminiInstrumentation
 from .mistral import MistralInstrumentation
+from .embedchain import EmbedchainInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
     "ChromaInstrumentation",
     "CohereInstrumentation",
     "CrewAIInstrumentation",
+    "EmbedchainInstrumentation",
     "GroqInstrumentation",
     "LangchainInstrumentation",
     "LangchainCommunityInstrumentation",
diff --git a/src/langtrace_python_sdk/instrumentation/chroma/patch.py b/src/langtrace_python_sdk/instrumentation/chroma/patch.py
@@ -51,7 +51,7 @@ def traced_method(wrapped, instance, args, kwargs):
             "langtrace.version": v(LANGTRACE_SDK_NAME),
             "db.system": "chromadb",
             "db.operation": api["OPERATION"],
-            "db.query": json.dumps(kwargs.get("query")),
+            "db.query": json.dumps(kwargs),
             **(extra_attributes if extra_attributes is not None else {}),
         }
 
diff --git a/src/langtrace_python_sdk/instrumentation/crewai/patch.py b/src/langtrace_python_sdk/instrumentation/crewai/patch.py
@@ -114,12 +114,23 @@ def traced_method(wrapped, instance, args, kwargs):
                 result = wrapped(*args, **kwargs)
                 if result:
                     span.set_status(Status(StatusCode.OK))
+                    if instance.__class__.__name__ == "Crew":
+                        span.set_attribute("crewai.crew.result", str(result))
+                        if hasattr(result, "tasks_output"):
+                            span.set_attribute("crewai.crew.tasks_output", str((result.tasks_output)))
+                        if hasattr(result, "token_usage"):
+                            span.set_attribute("crewai.crew.token_usage", str((result.token_usage)))
+                        if hasattr(result, "usage_metrics"):
+                            span.set_attribute("crewai.crew.usage_metrics", str((result.usage_metrics)))
+                    elif instance.__class__.__name__ == "Agent":
+                        span.set_attribute("crewai.agent.result", str(result))
+                    elif instance.__class__.__name__ == "Task":
+                        span.set_attribute("crewai.task.result", str(result))
 
                 span.end()
                 return result
 
             except Exception as err:
-                print("Error", err)
                 # Record the exception in the span
                 span.record_exception(err)
 
diff --git a/src/langtrace_python_sdk/instrumentation/embedchain/__init__.py b/src/langtrace_python_sdk/instrumentation/embedchain/__init__.py
@@ -0,0 +1,5 @@
+from .instrumentation import EmbedchainInstrumentation
+
+__all__ = [
+    "EmbedchainInstrumentation",
+]
diff --git a/src/langtrace_python_sdk/instrumentation/embedchain/instrumentation.py b/src/langtrace_python_sdk/instrumentation/embedchain/instrumentation.py
@@ -0,0 +1,65 @@
+"""
+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.
+"""
+
+import importlib.metadata
+import logging
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from langtrace_python_sdk.instrumentation.embedchain.patch import generic_patch
+
+logging.basicConfig(level=logging.FATAL)
+
+
+class EmbedchainInstrumentation(BaseInstrumentor):
+    """
+    The EmbedchainInstrumentation class represents the Embedchain instrumentation
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["embedchain >= 0.1.113"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version("embedchain")
+
+        wrap_function_wrapper(
+            "embedchain.embedchain",
+            "EmbedChain.add",
+            generic_patch("ADD", version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "embedchain.embedchain",
+            "EmbedChain.query",
+            generic_patch("QUERY", version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "embedchain.embedchain",
+            "EmbedChain.search",
+            generic_patch("SEARCH", version, tracer),
+        )
+
+    def _instrument_module(self, module_name):
+        pass
+
+    def _uninstrument(self, **kwargs):
+        pass
diff --git a/src/langtrace_python_sdk/instrumentation/embedchain/patch.py b/src/langtrace_python_sdk/instrumentation/embedchain/patch.py
@@ -0,0 +1,91 @@
+"""
+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 langtrace.trace_attributes import FrameworkSpanAttributes
+from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.llm import get_span_name
+from opentelemetry import baggage, trace
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.trace.propagation import set_span_in_context
+from langtrace_python_sdk.constants.instrumentation.embedchain import APIS
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
+import json
+from importlib_metadata import version as v
+
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+
+
+def generic_patch(method, version, tracer):
+    """
+    A generic patch method that wraps a function with a span
+    """
+
+    def traced_method(wrapped, instance, args, kwargs):
+        api = APIS[method]
+        service_provider = SERVICE_PROVIDERS["EMBEDCHAIN"]
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "framework",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "embedchain.api": api["OPERATION"],
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        if len(args) > 0:
+            span_attributes["embedchain.inputs"] = json.dumps(args)
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=get_span_name(api["METHOD"]),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            for field, value in attributes.model_dump(by_alias=True).items():
+                if value is not None:
+                    span.set_attribute(field, value)
+            try:
+                result = wrapped(*args, **kwargs)
+                set_span_attribute(span, "embedchain.outputs", json.dumps(result))
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as err:
+                # Record the exception in the span
+                span.record_exception(err)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def get_count_or_none(value):
+    return len(value) if value is not None else None
+
+
+def handle_null_params(param):
+    return str(param) if param else None
diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
@@ -42,6 +42,7 @@
     ChromaInstrumentation,
     CohereInstrumentation,
     CrewAIInstrumentation,
+    EmbedchainInstrumentation,
     GroqInstrumentation,
     LangchainInstrumentation,
     LangchainCommunityInstrumentation,
@@ -115,6 +116,7 @@ def init(
         "pinecone": PineconeInstrumentation(),
         "llamaindex": LlamaindexInstrumentation(),
         "chroma": ChromaInstrumentation(),
+        "embedchain": EmbedchainInstrumentation(),
         "qdrant": QdrantInstrumentation(),
         "langchain": LangchainInstrumentation(),
         "langchain_core": LangchainCoreInstrumentation(),
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "2.2.31"
+__version__ = "2.3.0"

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@ dev = [`
`55`	`55`	`"google-generativeai",`
`56`	`56`	`"google-cloud-aiplatform",`
`57`	`57`	`"mistralai",`
	`58`	`+ "embedchain",`
`58`	`59`	`]`
`59`	`60`
`60`	`61`	`test = ["pytest", "pytest-vcr", "pytest-asyncio"]`
Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@`
`30`	`30`	`"VERTEXAI": "VertexAI",`
`31`	`31`	`"GEMINI": "Gemini",`
`32`	`32`	`"MISTRAL": "Mistral",`
	`33`	`+ "EMBEDCHAIN": "Embedchain",`
`33`	`34`	`}`
`34`	`35`
`35`	`36`	`LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ def traced_method(wrapped, instance, args, kwargs):`
`51`	`51`	`"langtrace.version": v(LANGTRACE_SDK_NAME),`
`52`	`52`	`"db.system": "chromadb",`
`53`	`53`	`"db.operation": api["OPERATION"],`
`54`		`- "db.query": json.dumps(kwargs.get("query")),`
	`54`	`+ "db.query": json.dumps(kwargs),`
`55`	`55`	`**(extra_attributes if extra_attributes is not None else {}),`
`56`	`56`	`}`
`57`	`57`