inital logger interface

Author: rushilpatel0

src/codegen/agents/__init__.py

@@ -0,0 +1 @@
+

src/codegen/agents/code_agent.py

@@ -3,9 +3,10 @@
 from uuid import uuid4
 
 from langchain.tools import BaseTool
-from langchain_core.messages import AIMessage
 from langsmith import Client
-
+from langchain_core.messages import AIMessage
+from codegen.agents.loggers import  ExternalLogger
+from codegen.agents.tracer import MessageStreamTracer
 from codegen.extensions.langchain.agent import create_codebase_agent
 from codegen.extensions.langchain.utils.get_langsmith_url import find_and_print_langsmith_run_url
 
@@ -42,7 +43,7 @@ def __init__(self, codebase: "Codebase", model_provider: str = "anthropic", mode
         self.project_name = os.environ.get("LANGCHAIN_PROJECT", "RELACE")
         print(f"Using LangSmith project: {self.project_name}")
 
-    def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
+    def run(self, prompt: str, thread_id: Optional[str] = None, logger: Optional[ExternalLogger] = None) -> str:
         """Run the agent with a prompt.
 
         Args:
@@ -55,20 +56,32 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
         if thread_id is None:
             thread_id = str(uuid4())
 
+
+
         # this message has a reducer which appends the current message to the existing history
         # see more https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers
         input = {"messages": [("user", prompt)]}
 
         # we stream the steps instead of invoke because it allows us to access intermediate nodes
-        stream = self.agent.stream(input, config={"configurable": {"thread_id": thread_id, "metadata": {"project": self.project_name}}, "recursion_limit": 100}, stream_mode="values")
+        stream = self.agent.stream(input, config={"configurable": {"thread_id": thread_id,  "metadata": {"project": self.project_name}}, "recursion_limit": 100}, stream_mode="values")
+        # Get the stream from the graph
+                # Create the external logger (optional)        
+        # Create the tracer
+        _tracer = MessageStreamTracer(logger=logger)
+
+        # Process the stream with the tracer
+        traced_stream = _tracer.process_stream(stream)
 
         # Keep track of run IDs from the stream
         run_ids = []
 
-        for s in stream:
+        for s in traced_stream:
             message = s["messages"][-1]
+            message.pretty_print()
+
             if isinstance(message, tuple):
-                print(message)
+                # print(message)
+                pass
             else:
                 if isinstance(message, AIMessage) and isinstance(message.content, list) and "text" in message.content[0]:
                     AIMessage(message.content[0]["text"]).pretty_print()
@@ -82,7 +95,7 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
         # Get the last message content
         result = s["messages"][-1].content
 
-        # Try to find run IDs in the LangSmith client's recent runs
+        # # Try to find run IDs in the LangSmith client's recent runs
         try:
             # Find and print the LangSmith run URL
             find_and_print_langsmith_run_url(self.langsmith_client, self.project_name)

src/codegen/agents/data.py

@@ -0,0 +1,52 @@
+from typing import List, Optional
+from dataclasses import dataclass, field
+from datetime import datetime
+
+# Base dataclass for all message types
+@dataclass
+class BaseMessage:
+    """Base class for all message types."""
+    type: str
+    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
+    content: str = ""
+
+@dataclass
+class UserMessage(BaseMessage):
+    """Represents a message from the user."""
+    type: str = field(default="user")
+
+@dataclass
+class SystemMessageData(BaseMessage):
+    """Represents a system message."""
+    type: str = field(default="system")
+
+@dataclass
+class ToolCall:
+    """Represents a tool call within an assistant message."""
+    name: Optional[str] = None
+    arguments: Optional[str] = None
+    id: Optional[str] = None
+
+@dataclass
+class AssistantMessage(BaseMessage):
+    """Represents a message from the assistant."""
+    type: str = field(default="assistant")
+    tool_calls: List[ToolCall] = field(default_factory=list)
+
+@dataclass
+class ToolMessageData(BaseMessage):
+    """Represents a tool response message."""
+    type: str = field(default="tool")
+    tool_name: Optional[str] = None
+    tool_response: Optional[str] = None
+    tool_id: Optional[str] = None
+
+@dataclass
+class FunctionMessageData(BaseMessage):
+    """Represents a function message."""
+    type: str = field(default="function")
+
+@dataclass
+class UnknownMessage(BaseMessage):
+    """Represents an unknown message type."""
+    type: str = field(default="unknown") 

src/codegen/agents/loggers.py

@@ -0,0 +1,17 @@
+import json
+from typing import Dict, List, Any, Union, Protocol
+from dataclasses import asdict
+from .data import BaseMessage
+
+# Define the interface for ExternalLogger
+class ExternalLogger(Protocol):
+    """Protocol defining the interface for external loggers."""
+    
+    def log(self, data: Union[Dict[str, Any], BaseMessage]) -> None:
+        """
+        Log structured data to an external system.
+        
+        Args:
+            data: The structured data to log, either as a dictionary or a BaseMessage
+        """
+        pass