Merge pull request #814 from parea-ai/PAI-1075-trace-decorator-with-l…

…ang-chain-integration

feat: trace langchain calls within decorator

parea/cookbook/langchain/trace_langchain_inside_trace_decorator.py

@@ -0,0 +1,38 @@
+import os
+
+from dotenv import load_dotenv
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+from openai import OpenAI
+
+from parea import Parea, trace
+from parea.utils.trace_integrations.langchain import PareaAILangchainTracer
+
+load_dotenv()
+
+oai_client = OpenAI()
+
+p = Parea(api_key=os.getenv("PAREA_API_KEY"))
+handler = PareaAILangchainTracer()
+p.wrap_openai_client(oai_client)
+
+llm = ChatOpenAI(openai_api_key=os.getenv("OPENAI_API_KEY"))
+prompt = ChatPromptTemplate.from_messages([("user", "{input}")])
+chain = prompt | llm | StrOutputParser()
+
+
+@trace
+def main():
+    programming_language = (
+        oai_client.chat.completions.create(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Suggest one programming languages"}]).choices[0].message.content
+    )
+
+    return chain.invoke(
+        {"input": f"Write a Hello World program in {programming_language}."},
+        config={"callbacks": [handler]},
+    )
+
+
+if __name__ == "__main__":
+    print(main())

parea/schemas/models.py

@@ -220,6 +220,7 @@ class UpdateTraceScenario(str, Enum):
     CHAIN: str = "chain"
     USAGE: str = "usage"
     OPENAICONFIG: str = "openaiconfig"
+    LANGCHAIN_CHILD: str = "langchain_child"
 
 
 @define

parea/utils/trace_integrations/langchain.py

@@ -7,23 +7,38 @@
 
 from parea.constants import TURN_OFF_PAREA_LOGGING
 from parea.parea_logger import parea_logger
+from parea.schemas import UpdateTraceScenario
 from parea.schemas.log import TraceIntegrations
+from parea.utils.trace_utils import fill_trace_data, get_current_trace_id, get_root_trace_id
 
 
 class PareaAILangchainTracer(BaseTracer):
     parent_trace_id: UUID
+    _parea_root_trace_id: str = None
+    _parea_parent_trace_id: str = None
 
     def _persist_run(self, run: Union[Run, LLMRun, ChainRun, ToolRun]) -> None:
         if TURN_OFF_PAREA_LOGGING:
             return
         self.parent_trace_id = run.id
         # using .dict() since langchain Run class currently set to Pydantic v1
         data = run.dict()
+        data["_parea_root_trace_id"] = self._parea_root_trace_id or None
+        if run.execution_order == 1:
+            data["_parea_parent_trace_id"] = self._parea_parent_trace_id or None
         parea_logger.record_vendor_log(data, TraceIntegrations.LANGCHAIN)
 
     def get_parent_trace_id(self) -> UUID:
         return self.parent_trace_id
 
+    def _on_run_create(self, run: Run) -> None:
+        if run.execution_order == 1:
+            # need to check if any traces already exist\
+            self._parea_root_trace_id = get_root_trace_id()
+            if parent_trace_id := get_current_trace_id():
+                self._parea_parent_trace_id = parent_trace_id
+                fill_trace_data(str(run.id), {"parent_trace_id": parent_trace_id}, UpdateTraceScenario.LANGCHAIN_CHILD)
+
     def _on_llm_end(self, run: Run) -> None:
         self._persist_run(run)
 

parea/utils/trace_utils.py

@@ -117,6 +117,8 @@ def fill_trace_data(trace_id: str, data: Dict[str, Any], scenario: UpdateTraceSc
         elif scenario == UpdateTraceScenario.CHAIN:
             trace_data.get()[trace_id].parent_trace_id = data["parent_trace_id"]
             trace_data.get()[data["parent_trace_id"]].children.append(trace_id)
+        elif scenario == UpdateTraceScenario.LANGCHAIN_CHILD:
+            trace_data.get()[data["parent_trace_id"]].children.append(trace_id)
         elif scenario == UpdateTraceScenario.OPENAICONFIG:
             trace_data.get()[trace_id].configuration = data["configuration"]
             trace_data.get()[trace_id].output = data["output"]

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "parea-ai"
 packages = [{ include = "parea" }]
-version = "0.2.145"
+version = "0.2.146"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <[email protected]>"]