diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 51ea43d..4dcb54a 100644 --- a/src/utils/__init__.py +++ b/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(): diff --git a/src/utils/agents/__init__.py b/src/utils/agents/__init__.py index a3eaa99..7b002c4 100644 --- a/src/utils/agents/__init__.py +++ b/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("_"): @@ -25,7 +28,7 @@ for path in package_dir.glob("*.py"): module_name = path.stem full_module_name = f"{package_name}.{module_name}" - + try: module = importlib.import_module(full_module_name) @@ -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 - diff --git a/src/utils/agents/base.py b/src/utils/agents/base.py index 94d4ea0..b6c4729 100644 --- a/src/utils/agents/base.py +++ b/src/utils/agents/base.py @@ -1,7 +1,8 @@ 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, + 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) - diff --git a/src/utils/agents/chat.py b/src/utils/agents/chat.py index bc8711e..4be43eb 100644 --- a/src/utils/agents/chat.py +++ b/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) diff --git a/src/utils/agents/types.py b/src/utils/agents/types.py index 5baa9b6..1c1efdb 100644 --- a/src/utils/agents/types.py +++ b/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: diff --git a/src/utils/rag.py b/src/utils/rag.py index 8bcfd9b..1b390e6 100644 --- a/src/utils/rag.py +++ b/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