Core 1.0.4

bpm-ai-core/bpm_ai_core/tracing/langsmith.py

@@ -23,11 +23,11 @@ def __init__(self):
         self.context_run = None
 
     def start_llm_trace(
-        self,
-        llm: Any,
-        messages: List,
-        current_try: int,
-        tools: Optional[List] = None
+            self,
+            llm: Any,
+            messages: List,
+            current_try: int,
+            tools: Optional[List] = None
     ):
         from bpm_ai_core.util.openai import messages_to_openai_dicts, json_schema_to_openai_function
         inputs = {
@@ -36,64 +36,66 @@ def start_llm_trace(
             "temperature": llm.temperature,
             "current_try": current_try,
             "max_tries": llm.max_retries + 1,
-            "tools": [json_schema_to_openai_function(f.name, f.description, f.args_schema) for f in tools] if tools else None
+            "tools": [json_schema_to_openai_function(f.name, f.description, f.args_schema) for f in
+                      tools] if tools else None
         }
-        self.start_trace(
+        self._start_trace(
             name=llm.__class__.__name__,
             run_type="llm",
             inputs=inputs
         )
 
-    def end_llm_trace(self, completion = None, error_msg: Optional[str] = None):
+    def end_llm_trace(self, completion=None, error_msg: Optional[str] = None):
         from bpm_ai_core.llm.common.message import ToolCallsMessage
         choices = {
             "choices": [{
                 "message": {
                     "role": "assistant",
                     "content": completion.content,
                     **({"tool_calls":
-                            [{"function": {"name": c.name, "arguments": c.payload_dict()}} for c in completion.tool_calls]
+                            [{"function": {"name": c.name, "arguments": c.payload_dict()}} for c in
+                             completion.tool_calls]
                         } if isinstance(completion, ToolCallsMessage) else {})
                 }
             }]
         } if completion else None
-        self.end_trace(
+        self._end_trace(
             outputs=choices,
             error=error_msg
         )
 
     def start_function_trace(
-        self,
-        function,
-        inputs: dict,
+            self,
+            function,
+            inputs: dict,
     ):
-        self.start_trace(
+        self._start_trace(
             name=function.name,
             run_type="tool",
             inputs=inputs,
             metadata={"description": function.description}
         )
 
     def end_function_trace(
-        self,
-        output: Optional[dict],
-        error_msg: Optional[str] = None
+            self,
+            output: Optional[dict] = None,
+            error_msg: Optional[str] = None
     ):
-        self.end_trace(
+        self._end_trace(
             outputs=output,
             error=error_msg
         )
 
-    def start_trace(
-        self,
-        name: str,
-        run_type: str,
-        inputs: dict,
-        executor: Optional[futures.ThreadPoolExecutor] = None,
-        metadata: Optional[Mapping[str, Any]] = None,
-        tags: Optional[List[str]] = None,
-        client: Optional[client.Client] = None,
-        extra: Optional[Dict] = None
+    def _start_trace(
+            self,
+            name: str,
+            run_type: str,
+            inputs: dict,
+            executor: Optional[futures.ThreadPoolExecutor] = None,
+            metadata: Optional[Mapping[str, Any]] = None,
+            tags: Optional[List[str]] = None,
+            client: Optional[client.Client] = None,
+            extra: Optional[Dict] = None
     ):
         self.context_run = _PARENT_RUN_TREE.get()
         self.run_container = self._setup_run(
@@ -109,7 +111,7 @@ def start_trace(
         _PROJECT_NAME.set(self.run_container["project_name"])
         _PARENT_RUN_TREE.set(self.run_container["new_run"])
 
-    def end_trace(self, outputs: Optional[Dict[str, Any]] = None, error: Optional[str] = None):
+    def _end_trace(self, outputs: Optional[Dict[str, Any]] = None, error: Optional[str] = None):
         if self.run_container is None:
             raise RuntimeError("Must call start_trace() before end_trace()")
 
@@ -134,15 +136,15 @@ def end_trace(self, outputs: Optional[Dict[str, Any]] = None, error: Optional[st
 
     @staticmethod
     def _setup_run(
-        name: str,
-        run_type: str,
-        inputs: dict,
-        extra_outer: dict,
-        langsmith_extra: Optional[LangSmithExtra] = None,
-        executor: Optional[futures.ThreadPoolExecutor] = None,
-        metadata: Optional[Mapping[str, Any]] = None,
-        tags: Optional[List[str]] = None,
-        langsmith_client: Optional[client.Client] = None,
+            name: str,
+            run_type: str,
+            inputs: dict,
+            extra_outer: dict,
+            langsmith_extra: Optional[LangSmithExtra] = None,
+            executor: Optional[futures.ThreadPoolExecutor] = None,
+            metadata: Optional[Mapping[str, Any]] = None,
+            tags: Optional[List[str]] = None,
+            langsmith_client: Optional[client.Client] = None,
     ) -> _TraceableContainer:
         outer_project = _PROJECT_NAME.get() or os.environ.get(
             "LANGCHAIN_PROJECT", os.environ.get("LANGCHAIN_PROJECT", "default")
@@ -197,4 +199,4 @@ def _setup_run(
             outer_project=outer_project,
             outer_metadata=outer_metadata,
             outer_tags=outer_tags,
-        )
+        )

bpm-ai-core/bpm_ai_core/tracing/tracing.py

@@ -3,45 +3,31 @@
 
 class Tracer(metaclass=ABCMeta):
     @abstractmethod
-    def start_llm_trace(self, llm, messages, current_try, tools):
+    def start_llm_trace(self, llm, messages, current_try, tools=None):
         pass
 
     @abstractmethod
-    def end_llm_trace(self, completion, error_msg):
+    def end_llm_trace(self, completion=None, error_msg=None):
         pass
 
     @abstractmethod
     def start_function_trace(self, function, inputs):
         pass
 
     @abstractmethod
-    def end_function_trace(self, output, error_msg):
-        pass
-
-    @abstractmethod
-    def start_trace(self, name, run_type, inputs, executor, metadata, tags, client, extra):
-        pass
-
-    @abstractmethod
-    def end_trace(self, outputs, error):
+    def end_function_trace(self, output=None, error_msg=None):
         pass
 
 
 class NoopTracer(Tracer):
-    def start_llm_trace(self, llm, messages, current_try, tools):
+    def start_llm_trace(self, llm, messages, current_try, tools=None):
         pass
 
-    def end_llm_trace(self, completion, error_msg):
+    def end_llm_trace(self, completion=None, error_msg=None):
         pass
 
     def start_function_trace(self, function, inputs):
         pass
 
-    def end_function_trace(self, output, error_msg):
-        pass
-
-    def start_trace(self, name, run_type, inputs, executor, metadata, tags, client, extra):
+    def end_function_trace(self, output=None, error_msg=None):
         pass
-
-    def end_trace(self, outputs, error):
-        pass

bpm-ai-core/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bpm-ai-core"
-version = "1.0.3"
+version = "1.0.4"
 description = "Core AI abstractions and helpers."
 authors = ["Bennet Krause <[email protected]>"]
 repository = "https://github.com/holunda-io/bpm-ai"