Cleaning up import orders

src/utils/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
-
-import importlib
 from pydantic import BaseModel # type: ignore
+import importlib
 
 from . import defines  
 from . context import Context
@@ -9,7 +8,7 @@ from . conversation import Conversation
 from . message import Message
 from . rag import ChromaDBFileWatcher, start_file_watcher
 from . setup_logging import setup_logging
-from .agents import class_registry, AnyAgent, Agent, __all__ as agents_all
+from . agents import class_registry, AnyAgent, Agent, __all__ as agents_all
 
 __all__ = [
   'Agent', 
@@ -23,13 +22,6 @@ __all__ = [
 
 __all__.extend(agents_all) # type: ignore
 
-# Resolve circular dependencies by rebuilding models
-# Call model_rebuild() on Agent and Context
-Agent.model_rebuild()
-Context.model_rebuild()
-
-# Assuming class_registry is available from agents/__init__.py
-
 logger = setup_logging(level=defines.logging_level)
 
 def rebuild_models():

src/utils/agents/__init__.py

@@ -1,23 +1,26 @@
 from __future__ import annotations
+from typing import TypeAlias, Dict, Tuple
 import importlib
 import pathlib
 import inspect
-from typing import TypeAlias, Dict, Tuple
+
 from . types import registry
-from . base import Agent
 from .. setup_logging import setup_logging
 from .. import defines
+from . base import Agent
+
 
 logger = setup_logging(defines.logging_level)
 
+__all__ = [ "AnyAgent", "registry", "class_registry" ]
+
 # Type alias for Agent or any subclass
 AnyAgent: TypeAlias = Agent  # BaseModel covers Agent and subclasses
 
+class_registry: Dict[str, Tuple[str, str]] = {}  # Maps class_name to (module_name, class_name)
+
 package_dir = pathlib.Path(__file__).parent
 package_name = __name__
-__all__ = [ "AnyAgent", "registry"]
-
-class_registry: Dict[str, Tuple[str, str]] = {}  # Maps class_name to (module_name, class_name)
 
 for path in package_dir.glob("*.py"):
     if path.name in ("__init__.py", "base.py") or path.name.startswith("_"):
@@ -47,4 +50,3 @@ for path in package_dir.glob("*.py"):
     except Exception as e:
         logger.error(f"Error processing {full_module_name}: {e}")
         raise e
-

src/utils/agents/base.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 from pydantic import BaseModel, PrivateAttr, Field # type: ignore
 from typing import (
     Literal, get_args, List, AsyncGenerator, TYPE_CHECKING, Optional, ClassVar, Any,
+    TypeAlias, Dict, Tuple
 )
 from abc import ABC
 from .. setup_logging import setup_logging
@@ -109,6 +110,57 @@ class Agent(BaseModel, ABC):
     def get_agent_type(self):
         return self._agent_type
     
+    async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
+        """
+        Prepare message with context information in message.preamble
+        """
+        logging.info(f"{self.agent_type} - {inspect.stack()[1].function}")
+
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
+        
+        # Generate RAG content if enabled, based on the content
+        rag_context = ""
+        if message.enable_rag:
+            # Gather RAG results, yielding each result
+            # as it becomes available
+            for message in self.context.generate_rag_results(message):
+                logging.info(f"RAG: {message.status} - {message.response}")
+                if message.status == "error":
+                    yield message
+                    return
+                if message.status != "done":
+                    yield message
+                
+            if "rag" in message.metadata and message.metadata["rag"]:
+                for rag in message.metadata["rag"]:
+                    for doc in rag["documents"]:
+                        rag_context += f"{doc}\n"
+
+        message.preamble = {}
+
+        if rag_context:
+            message.preamble["context"] = rag_context
+
+        if self.context.user_resume:
+            message.preamble["resume"] = self.context.user_resume
+
+        if message.preamble:
+            preamble_types = [f"<|{p}|>" for p in message.preamble.keys()]
+            preamble_types_AND = " and ".join(preamble_types)
+            preamble_types_OR = " or ".join(preamble_types)
+            message.preamble["rules"] = f"""\
+    - Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
+    - If there is no information in these sections, answer based on your knowledge, or use any available tools.
+    - Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
+    """
+        message.preamble["question"] = "Respond to:"
+        
+        message.system_prompt = self.system_prompt
+        message.status = "done"
+        yield message
+        return
+
     async def process_tool_calls(self, llm: Any, model: str, message: Message, tool_message: Any, messages: List[Any]) -> AsyncGenerator[Message, None]:
         logging.info(f"{self.agent_type} - {inspect.stack()[1].function}")
 
@@ -422,4 +474,3 @@ class Agent(BaseModel, ABC):
 # Register the base agent
 registry.register(Agent._agent_type, Agent)
 
-

src/utils/agents/chat.py

@@ -12,56 +12,5 @@ class Chat(Agent):
   agent_type: Literal["chat"] = "chat" # type: ignore
   _agent_type: ClassVar[str] = agent_type  # Add this for registration
 
-  async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
-    """
-    Prepare message with context information in message.preamble
-    """
-    logging.info(f"{self.agent_type} - {inspect.stack()[1].function}")
-
-    if not self.context:
-       raise ValueError("Context is not set for this agent.")
-    
-    # Generate RAG content if enabled, based on the content
-    rag_context = ""
-    if message.enable_rag:
-        # Gather RAG results, yielding each result
-        # as it becomes available
-        for message in self.context.generate_rag_results(message):
-            logging.info(f"RAG: {message.status} - {message.response}")
-            if message.status == "error":
-                yield message
-                return
-            if message.status != "done":
-                yield message
-            
-        if "rag" in message.metadata and message.metadata["rag"]:
-            for rag in message.metadata["rag"]:
-                for doc in rag["documents"]:
-                  rag_context += f"{doc}\n"
-
-    message.preamble = {}
-
-    if rag_context:
-        message.preamble["context"] = rag_context
-
-    if self.context.user_resume:
-        message.preamble["resume"] = self.context.user_resume
-
-    if message.preamble:
-        preamble_types = [f"<|{p}|>" for p in message.preamble.keys()]
-        preamble_types_AND = " and ".join(preamble_types)
-        preamble_types_OR = " or ".join(preamble_types)
-        message.preamble["rules"] = f"""\
-- Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
-- If there is no information in these sections, answer based on your knowledge, or use any available tools.
-- Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
-"""
-    message.preamble["question"] = "Respond to:"
-    
-    message.system_prompt = self.system_prompt
-    message.status = "done"
-    yield message
-    return
-
 # Register the base agent
 registry.register(Chat._agent_type, Chat)

src/utils/agents/types.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 from typing import List, Dict, ForwardRef, Optional, Type
 
 # Forward references
-AgentRef = ForwardRef('Agent')
-ContextRef = ForwardRef('Context')
+# AgentRef = ForwardRef('Agent')
+# ContextRef = ForwardRef('Context')
 
 # We'll use a registry pattern rather than hardcoded strings
 class AgentRegistry:

src/utils/rag.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, # type: ignore
+from pydantic import BaseModel # type: ignore
 from typing import List, Optional, Dict, Any
 import os
 import glob