llumo 0.2.28__py3-none-any.whl → 0.2.30__py3-none-any.whl

llumo/callback.py

@@ -11,11 +11,12 @@ import re
 
 
 class LlumoCallbackHandler(BaseCallbackHandler):
-    def __init__(self, session: LlumoSessionContext = None):
+    def __init__(self, session: LlumoSessionContext = None,agentType = "react_agent"):
         if session is None:
             raise ValueError("LlumoSessionContext is required")
 
         self.sessionLogger = session
+        self.agentType = agentType
 
         # Initialize timing and state variables
         self.llmStartTime = None
@@ -26,6 +27,7 @@ class LlumoCallbackHandler(BaseCallbackHandler):
 
         # Initialize tracking variables
         self.prompt = ""
+        self.searchQuery = ""
         self.currentToolName = None
         self.currentToolInput = None
         self.currentAgentName = None
@@ -45,8 +47,14 @@ class LlumoCallbackHandler(BaseCallbackHandler):
         self.currentObservation = ""
         self.isAgentExecution = False
 
+
+
     def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
         """Called when a chain starts - this includes agent execution"""
+        # print("ON CHAIN START: ",inputs)
+        # print("ON CHAIN START: serialized",serialized)
+        # print("ON CHAIN START: kwargs",kwargs)
+
         try:
             self.prompt = inputs.get("input", "")
             self.chainStartTime = time.time()
@@ -62,6 +70,7 @@ class LlumoCallbackHandler(BaseCallbackHandler):
             self.currentToolInput = None
             self.hasErrors = False
             self.lastError = None
+            self.toolDescription = {}
 
             # Dynamically detect agent name from serialized data
             if serialized is not None:
@@ -91,6 +100,26 @@ class LlumoCallbackHandler(BaseCallbackHandler):
         except Exception as e:
             print(f"[ERROR] in on_chain_start: {e}")
 
+        try:
+            self.sessionLogger.logQueryStep(
+                stepName = "Query Invocation",
+                model = "unknown",
+                provider = "unknown",
+                inputTokens = round(len(self.prompt.split()) * 1.5),
+                query = self.prompt,
+                status = "SUCCESS"
+            )
+        except Exception as e:
+            self.sessionLogger.logQueryStep(
+                stepName="Query Invocation",
+                model="unknown",
+                provider="unknown",
+                inputTokens=0,
+                query="",
+                status="FAILURE"
+            )
+            print(f"[ERROR] Failed to log user input: {e}")
+
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Called when a chain ends"""
         try:
@@ -98,15 +127,16 @@ class LlumoCallbackHandler(BaseCallbackHandler):
                 # Use logAgentStep for final completion
                 self.sessionLogger.logAgentStep(
                     stepName="Agent Execution Completed",
-                    agentType="react_agent",
+                    agentType=self.agentType,
                     agentName=self.currentAgentName or "unknown",
                     numStepsTaken=self.agentsSteps,
                     tools=self.toolsUsed,
                     query=self.prompt,
                     status="SUCCESS",
-                    message=f"Final output: {outputs['output']}. ReAct steps: {json.dumps(self.reactSteps)}",
+                    # message=f"Final output: {outputs['output']}. ReAct steps: {json.dumps(self.reactSteps)}",
                 )
 
+
             # Reset execution state after chain ends
             self.isAgentExecution = False
 
@@ -118,6 +148,14 @@ class LlumoCallbackHandler(BaseCallbackHandler):
         self.llmStartTime = time.time()
         self.stepTime = time.time()
 
+        if self.prompt == "":
+            match = re.search(r"Human:\s*(.*)",prompts[0], re.DOTALL)
+            if match:
+                user_question = match.group(1).strip()
+                self.prompt = user_question # 👉 What is LangChain?
+            else:
+                self.prompt = ""
+
         # Dynamically get model info
         model = "unknown"
         if serialized and "kwargs" in serialized:
@@ -204,7 +242,6 @@ class LlumoCallbackHandler(BaseCallbackHandler):
         # Parse ReAct reasoning from LLM output if we're in agent execution
         if self.isAgentExecution and output:
             self._parse_react_reasoning(output)
-
         try:
             self.sessionLogger.logLlmStep(
                 stepName="LLM Call Completed",
@@ -212,17 +249,62 @@ class LlumoCallbackHandler(BaseCallbackHandler):
                 provider=self.llmProvider,
                 inputTokens=int(input_tokens),
                 outputTokens=int(output_tokens),
-                temperature=float(kwargs.get("temperature", 0.7)),
-                promptTruncated=False,
+                # temperature=float(kwargs.get("temperature", 0.7)),
+                # promptTruncated=False,
                 latencyMs=duration_ms,
-                query=str(self.prompt),
+                prompt=str(self.prompt),
                 output=output,
                 status=status,
-                message=error_message if status == "ERROR" else "",
+                # message=error_message if status == "ERROR" else "",
             )
+
         except Exception as e:
             print(f"[ERROR] Failed to log LLM end: {e}")
 
+    def on_retriever_start(self, serialized, query, run_id, parent_run_id=None, **kwargs):
+        self.prompt = query
+        self.searchQuery = query
+
+    def on_retriever_end(self, documents, run_id, parent_run_id=None, **kwargs):
+
+        try:
+            chunkSize = len(documents[0].page_content) if documents and documents[0].page_content else 0
+        except Exception:
+            chunkSize = 0
+
+        source = ( kwargs.get("metadata", {}).get("source") or kwargs.get("tags") or "unknown")
+
+        try:
+            self.sessionLogger.logRetrieverStep(
+                stepName="Context Retrieval Complete",
+                retrieverSource = str(source),
+                topK = len(documents),
+                chunkSize =  chunkSize,
+                context =  [doc.page_content for doc in documents],
+                searchQuery =  self.prompt if self.prompt != "" else self.searchQuery,
+                latencyMs = 120,  # mock latency, replace with real timing if needed
+                status = "SUCCESS"
+             )
+        except Exception as e:
+            print(f"[ERROR] Failed to log chain output: {e}")
+
+    def on_retriever_error(self, error, run_id, parent_run_id=None, **kwargs):
+
+        try:
+            self.sessionLogger.logRetrieverStep(
+                stepName="Context Retrieval Error",
+                retrieverSource = kwargs.get("metadata", {}).get("source", "unknown"),
+                topK = 0,
+                chunkSize =  0,
+                context = [],
+                searchQuery =  self.prompt if self.prompt != "" else self.searchQuery,
+                latencyMs = 0,  # mock latency, replace with real timing if needed
+                status = "FAILURE"
+             )
+        except Exception as e:
+            print(f"[ERROR] Failed to log chain output: {e}")
+
+
     def _parse_react_reasoning(self, llm_output: str):
         """Parse ReAct reasoning pattern from LLM output"""
         try:
@@ -265,6 +347,9 @@ class LlumoCallbackHandler(BaseCallbackHandler):
 
     def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None:
         """Called when a tool starts executing"""
+        # print("ON TOOL START: ",serialized)
+        # print("ON TOOL START: ",kwargs)
+
         self.toolStartTime = time.time()
         self.stepTime = time.time()
 
@@ -272,6 +357,7 @@ class LlumoCallbackHandler(BaseCallbackHandler):
         self.currentToolName = (serialized.get("name") or
                                 serialized.get("_type") or
                                 "unknown")
+        self.currentToolDescription = serialized.get("description","No description found")
 
         # Handle the case where input_str is "None" or None
         if input_str == "None" or input_str is None:
@@ -316,13 +402,13 @@ class LlumoCallbackHandler(BaseCallbackHandler):
             self.sessionLogger.logToolStep(
                 stepName="Tool Execution Completed",
                 toolName=self.currentToolName or "unknown",
+                description = self.currentToolDescription,
                 input=self.currentToolInput or {"input": ""},
                 output=output_str,
                 latencyMs=duration_ms,
                 status="SUCCESS",
-                message="",
+                # message="",
             )
-
             print(f"[DEBUG] Tool completed: {self.currentToolName} -> {output_str}")
 
         except Exception as e:
@@ -331,7 +417,7 @@ class LlumoCallbackHandler(BaseCallbackHandler):
     def on_agent_action(self, action: AgentAction, **kwargs: Any) -> None:
         """Called when an agent takes an action"""
         self.agentsSteps += 1
-        print("ON AGENT ACTION: ", action)
+        # print("ON AGENT ACTION: ", action)
 
         try:
             # Dynamically extract information from action
@@ -352,52 +438,54 @@ class LlumoCallbackHandler(BaseCallbackHandler):
                 })
 
             # Log the agent action step using logAgentStep
-            current_status = "ERROR" if self.hasErrors else "SUCCESS"
+            current_status = "FAILURE" if self.hasErrors else "SUCCESS"
             reasoning_text = self.currentThought if self.currentThought else "No reasoning captured"
 
             self.sessionLogger.logAgentStep(
                 stepName=f"Agent Action Step {self.agentsSteps}",
-                agentType="react_agent",
+                agentType=self.agentType,
                 agentName=self.currentAgentName or "unknown",
                 numStepsTaken=self.agentsSteps,
                 tools=[tool_name],
                 query=self.prompt,
                 status=current_status,
-                message=f"Executing {tool_name} with input: {tool_input}. Reasoning: {reasoning_text}",
+                # message=f"Executing {tool_name} with input: {tool_input}. Reasoning: {reasoning_text}",
             )
 
+
+
         except Exception as e:
             print(f"[ERROR] Failed to log agent action: {e}")
 
     def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
         """Called when an agent completes execution"""
-        print("ON AGENT FINISH:", finish)
+        # print("ON AGENT FINISH:", finish)
         # We don't need to log anything here since the final result is already logged in on_chain_end
         pass
 
     def on_agent_error(self, error: Exception, **kwargs: Any) -> None:
         """Called when an agent encounters an error"""
-        print("ITS A AGENT ERROR:", error)
+        # print("ITS A AGENT ERROR:", error)
         self.hasErrors = True
         self.lastError = str(error)
 
         try:
             self.sessionLogger.logAgentStep(
                 stepName="Agent Execution Error",
-                agentType="react_agent",
+                agentType=self.agentType,
                 agentName=self.currentAgentName or "unknown",
                 numStepsTaken=self.agentsSteps,
                 tools=self.toolsUsed,
                 query=self.prompt,
-                status="ERROR",
-                message=str(error),
+                status="FAILURE",
+                # message=str(error),
             )
         except Exception as e:
             print(f"[ERROR] Failed to log agent error: {e}")
 
     def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
         """Called when a tool encounters an error"""
-        print("ITS A TOOL ERROR:", error)
+
         self.hasErrors = True
         self.lastError = str(error)
 
@@ -410,18 +498,19 @@ class LlumoCallbackHandler(BaseCallbackHandler):
             self.sessionLogger.logToolStep(
                 stepName="Tool Execution Failed",
                 toolName=self.currentToolName or "unknown",
+                description=self.currentToolDescription,
                 input=self.currentToolInput or {"input": ""},
                 output="",
                 latencyMs=0,
-                status="ERROR",
-                message=str(error),
+                status="FAILURE",
+                # message=str(error),
             )
         except Exception as e:
             print(f"[ERROR] Failed to log tool error: {e}")
 
     def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
         """Called when a chain encounters an error"""
-        print("ITS A CHAIN ERROR:", error)
+        # print("ITS A CHAIN ERROR:", error)
         self.hasErrors = True
         self.lastError = str(error)
 
@@ -430,14 +519,17 @@ class LlumoCallbackHandler(BaseCallbackHandler):
                 # Use logAgentStep for agent-related chain errors
                 self.sessionLogger.logAgentStep(
                     stepName="Agent Chain Error",
-                    agentType="react_agent",
+                    agentType=self.agentType,
                     agentName=self.currentAgentName or "unknown",
                     numStepsTaken=self.agentsSteps,
                     tools=self.toolsUsed,
                     query=self.prompt,
-                    status="ERROR",
-                    message=str(error),
+                    status="FAILURE",
+                    # message=str(error),
                 )
+
+
+
             else:
                 # Use logLlmStep for general chain errors
                 self.sessionLogger.logLlmStep(
@@ -446,13 +538,13 @@ class LlumoCallbackHandler(BaseCallbackHandler):
                     provider=self.llmProvider,
                     inputTokens=0,
                     outputTokens=0,
-                    temperature=0.0,
-                    promptTruncated=False,
+                    # temperature=0.0,
+                    # promptTruncated=False,
                     latencyMs=0,
-                    query=self.prompt,
+                    prompt=self.prompt,
                     output="",
-                    status="ERROR",
-                    message=str(error),
+                    status="FAILURE",
+                    # message=str(error),
                 )
         except Exception as e:
             print(f"[ERROR] Failed to log chain error: {e}")