crewplus 0.2.43__py3-none-any.whl → 0.2.46__py3-none-any.whl

crewplus/callbacks/async_langfuse_handler.py

@@ -2,20 +2,29 @@
 import asyncio
 import contextvars
 from contextlib import contextmanager
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Optional, Sequence
+from uuid import UUID
 
 try:
     from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
     from langchain_core.callbacks import AsyncCallbackHandler
-    from langchain_core.outputs import LLMResult
+    from langchain_core.outputs import LLMResult, ChatGeneration
+    from langchain_core.messages import BaseMessage
+    from langchain.schema.agent import AgentAction, AgentFinish
+    from langchain.schema.document import Document
     LANGFUSE_AVAILABLE = True
 except ImportError:
     LANGFUSE_AVAILABLE = False
     LangfuseCallbackHandler = None
     AsyncCallbackHandler = object
+    # Define dummy types if langchain is not available
+    LLMResult = object
+    BaseMessage = object
+    AgentAction = object
+    AgentFinish = object
+    Document = object
+
 
-# This token is a simple flag to indicate that we are in an async context.
-# We use a context variable to make it available only within the async task.
 _ASYNC_CONTEXT_TOKEN = "in_async_context"
 in_async_context = contextvars.ContextVar(_ASYNC_CONTEXT_TOKEN, default=False)
 
@@ -30,13 +39,8 @@ def async_context():
 
 class AsyncLangfuseCallbackHandler(AsyncCallbackHandler):
     """
-    Wraps the synchronous LangfuseCallbackHandler to make it compatible with
-    LangChain's async methods.
-    
-    This works by running the synchronous handler's methods in a separate thread
-    using `asyncio.to_thread`. This is crucial because `asyncio`'s default
-    executor can correctly propagate `contextvars`, which solves the
-    `ValueError: <Token ...> was created in a different Context` from OpenTelemetry.
+    Wraps the synchronous LangfuseCallbackHandler to make it fully compatible with
+    LangChain's async methods by handling all relevant events.
     """
     def __init__(self, *args: Any, **kwargs: Any):
         if not LANGFUSE_AVAILABLE:
@@ -44,64 +48,115 @@ class AsyncLangfuseCallbackHandler(AsyncCallbackHandler):
         self.sync_handler = LangfuseCallbackHandler(*args, **kwargs)
 
     def __getattr__(self, name: str) -> Any:
-        # Delegate any other attribute access to the sync handler
         return getattr(self.sync_handler, name)
 
+    # LLM Events
     async def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+        self, serialized: Dict[str, Any], prompts: List[str], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
     ) -> None:
-        # --- DEBUGGING: Inspect the arguments from LangChain ---
-        print("--- [DEBUG] AsyncLangfuseCallbackHandler.on_llm_start ---")
-        print(f"Received prompts type: {type(prompts)}")
-        print(f"Received prompts value: {prompts!r}") # Using !r to see quotes if it's a string
-        print("----------------------------------------------------------")
-        # --- END DEBUGGING ---
-
-        # WORKAROUND: LangChain's async implementation can sometimes pass a raw
-        # string for prompts instead of a list. We wrap it in a list to ensure
-        # compatibility with the synchronous handler.
         corrected_prompts = prompts if isinstance(prompts, list) else [prompts]
-        
         await asyncio.to_thread(
-            self.sync_handler.on_llm_start, serialized, corrected_prompts, **kwargs
+            self.sync_handler.on_llm_start, serialized, corrected_prompts, run_id=run_id, parent_run_id=parent_run_id, **kwargs
+        )
+
+    async def on_chat_model_start(
+        self, serialized: Dict[str, Any], messages: List[List[BaseMessage]], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_chat_model_start, serialized, messages, run_id=run_id, parent_run_id=parent_run_id, **kwargs
+        )
+
+    async def on_llm_end(
+        self, response: LLMResult, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> None:
+        await asyncio.to_thread(
+            self.sync_handler.on_llm_end, response, run_id=run_id, parent_run_id=parent_run_id, **kwargs
+        )
+
+    async def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> None:
+        await asyncio.to_thread(
+            self.sync_handler.on_llm_error, error, run_id=run_id, parent_run_id=parent_run_id, **kwargs
+        )
+
+    # Chain Events
+    async def on_chain_start(
+        self, serialized: Dict[str, Any], inputs: Dict[str, Any], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_chain_start, serialized, inputs, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
 
-    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+    async def on_chain_end(
+        self, outputs: Dict[str, Any], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_llm_end, response, **kwargs
+            self.sync_handler.on_chain_end, outputs, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
-    
-    async def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
+
+    async def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_llm_error, error, **kwargs
+            self.sync_handler.on_chain_error, error, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
-    
-    async def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
+
+    # Tool Events
+    async def on_tool_start(
+        self, serialized: Dict[str, Any], input_str: str, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_tool_start, serialized, input_str, **kwargs
+            self.sync_handler.on_tool_start, serialized, input_str, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
 
-    async def on_tool_end(self, output: str, **kwargs: Any) -> Any:
+    async def on_tool_end(
+        self, output: str, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_tool_end, output, **kwargs
+            self.sync_handler.on_tool_end, output, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
 
-    async def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    async def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_tool_error, error, **kwargs
+            self.sync_handler.on_tool_error, error, run_id=run_id, parent_run_id=parent_run_id, **kwargs
+        )
+        
+    # Retriever Events
+    async def on_retriever_start(
+        self, serialized: Dict[str, Any], query: str, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_retriever_start, serialized, query, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
 
-    async def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> Any:
+    async def on_retriever_end(
+        self, documents: Sequence[Document], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_chain_start, serialized, inputs, **kwargs
+            self.sync_handler.on_retriever_end, documents, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
 
-    async def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
+    async def on_retriever_error(
+        self, error: Union[Exception, KeyboardInterrupt], *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
+        await asyncio.to_thread(
+            self.sync_handler.on_retriever_error, error, run_id=run_id, parent_run_id=parent_run_id, **kwargs
+        )
+        
+    # Agent Events
+    async def on_agent_action(
+        self, action: AgentAction, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_chain_end, outputs, **kwargs
+            self.sync_handler.on_agent_action, action, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )
 
-    async def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    async def on_agent_finish(
+        self, finish: AgentFinish, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any
+    ) -> Any:
         await asyncio.to_thread(
-            self.sync_handler.on_chain_error, error, **kwargs
+            self.sync_handler.on_agent_finish, finish, run_id=run_id, parent_run_id=parent_run_id, **kwargs
         )

{crewplus-0.2.43.dist-info → crewplus-0.2.46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.43
+Version: 0.2.46
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT

{crewplus-0.2.43.dist-info → crewplus-0.2.46.dist-info}/RECORD

@@ -1,10 +1,10 @@
-crewplus-0.2.43.dist-info/METADATA,sha256=h8suioxG3jmRN4AlcxKmcjuWByJHZaXrrJ3moEQYjO0,5362
-crewplus-0.2.43.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.43.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.43.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.46.dist-info/METADATA,sha256=QSaSV8WkRmj_qHybxSCTNlyESu9UNuiYNxeLgolNBAA,5362
+crewplus-0.2.46.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.46.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.46.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/callbacks/__init__.py,sha256=YG7ieeb91qEjp1zF0-inEN7mjZ7yT_D2yzdWFT8Z1Ws,63
-crewplus/callbacks/async_langfuse_handler.py,sha256=XnHpQ-sqcUCb6BCrE_JgldtTbaV61lH9aJnfy_IUp9Y,4614
+crewplus/callbacks/async_langfuse_handler.py,sha256=A4uFeLpvOUdc58M7sZoE65_C1V98u0QCvx5jUquM0pM,7006
 crewplus/services/__init__.py,sha256=V1CG8b2NOmRzNgQH7BPl4KVxWSYJH5vfEsW1wVErKNE,375
 crewplus/services/azure_chat_model.py,sha256=iWzJ2GQFSNmwJx-2O5_xKPSB6VVc-7T6bcfFI8_WezA,5521
 crewplus/services/gemini_chat_model.py,sha256=VsOB_st1qRmDkwLXzo-gCShhUsZHpk0V-G-ulQXGN3g,40081
@@ -22,4 +22,4 @@ docs/GeminiChatModel.md,sha256=zZYyl6RmjZTUsKxxMiC9O4yV70MC4TD-IGUmWhIDBKA,8677
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.43.dist-info/RECORD,,
+crewplus-0.2.46.dist-info/RECORD,,