some minor adjustments

Author: sha1cybr

run_ui.sh

@@ -1,2 +1,2 @@
 #!/bin/bash
-poetry run python src/agentwatch/cli.py ui
+AGENTWATCH_INTERNAL=1 PYTHONPATH=$PYTHONPATH:$(pwd)/src poetry run python src/agentwatch/cli.py ui

src/agentwatch/event_processor.py

@@ -25,8 +25,6 @@ class EventProcessor:
     NUM_WORKERS = 1
 
     def __init__(self) -> None:
-        logger.info("agentwatch initializing...")
-
         self._exit_event: Optional[Event] = None
         self._init_event: Optional[Event] = None
         self._pipe: Optional[Connection] = None
@@ -40,8 +38,6 @@ def __init__(self) -> None:
         self._supported_processors: list[type[BaseProcessor]] = [
             HttpProcessor
             ]
-        
-        logger.debug("agentwatch initialized!")
 
     def start(self, pipe: Connection, init_event: Event, exit_event: Event) -> None: 
         self._pipe = pipe
@@ -65,7 +61,7 @@ async def _start(self) -> None:
 
         self._register_visualization_webhook()
 
-        logger.debug(f"agentwatch up and running")
+        logger.info(f"agentwatch up and running!")
         if self._init_event:
             self._init_event.set()
 

src/agentwatch/graph/models.py

@@ -1,7 +1,7 @@
 import time
 import uuid
 from abc import abstractmethod
-from typing import Any, Type, TypeAlias
+from typing import Any, Optional, Type, TypeAlias
 
 from pydantic import BaseModel, Field
 
@@ -41,9 +41,13 @@ class ModelGenerateEdge(Edge):
     edge_type: EdgeType = EdgeType.MODEL_GENERATE
     prompt: str
 
+    class Config:
+        extra = "allow"
+
 class ToolCallEdge(Edge):
     edge_type: EdgeType = EdgeType.TOOL_CALL
     tool_input: dict[str, Any]
+    tool_name: Optional[str] = None
 
 GraphStructure: TypeAlias = tuple[list[Node], list[Edge]]
 

src/agentwatch/llm/anthropic_models.py

@@ -36,7 +36,8 @@ def extract_graph_structure(self) -> GraphStructure:
                     if isinstance(content, ToolUse):
                         tool_call_edge = ToolCallEdge(source_node_id=APP_NODE_ID,
                                                         target_node_id=content.name,
-                                                        tool_input=content.input)
+                                                        tool_input=content.input,
+                                                        tool_name=content.name)
                         edges.append(tool_call_edge)
                     elif isinstance(content, TextContent):
                         model_generate_edge = ModelGenerateEdge(prompt=content.text,

src/agentwatch/llm/openai_models.py

@@ -50,19 +50,16 @@ def extract_graph_structure(self) -> GraphStructure:
         for message in self.messages:
             if isinstance(message, ToolMessage):
                 pass
-                # Not sure I want this
+                # Not sure I want this, it includes the actual tool call result from the app
                 # edges.append(ToolCallEdge(source_node_id=APP_NODE_ID,
                 #                           target_node_id=self.model,
                 #                           tool_input={"input": message.content}))
             elif isinstance(message, UserMessage):
                 edges.append(ModelGenerateEdge(source_node_id=APP_NODE_ID,
                                               target_node_id=self.model,
-                                              prompt=message.content))
+                                              prompt=message.content,
+                                              history_size=len(self.messages)))  # type: ignore
             elif isinstance(message, AssistantMessage):
-                for tool_call in message.tool_calls:
-                    edges.append(ToolCallEdge(source_node_id=self.model,
-                                              target_node_id=APP_NODE_ID,
-                                              tool_input=json.loads(tool_call.function.arguments)))
                 if message.content is not None:
                     edges.append(ModelGenerateEdge(source_node_id=self.model,
                                                  target_node_id=APP_NODE_ID,
@@ -113,6 +110,14 @@ def extract_graph_structure(self) -> GraphStructure:
 
         for choice in self.choices:
             for tool_call in choice.message.tool_calls:
+                # Two edges should appear here: one edge displays the result from the LLM
+                # The other one 'simulates' the app call to the tool itself
+                edges.append(ToolCallEdge(source_node_id=self.model,
+                            target_node_id=APP_NODE_ID, 
+                            tool_name=tool_call.function.name,
+                            tool_input=json.loads(tool_call.function.arguments))
+                            )
+                
                 edges.append(ToolCallEdge(source_node_id=APP_NODE_ID, 
                                           target_node_id=tool_call.function.name, 
                                           tool_input=json.loads(tool_call.function.arguments)))

src/agentwatch/visualization/frontend/components/Sidebar.tsx

@@ -9,8 +9,6 @@ type EdgeData = Edge & {
   source: string;
   target: string;
   data?: {
-    prompt?: string;
-    tool_input?: string;
     [key: string]: any;
   };
 };