PR Summary: Refactor Embedding Handling and Environment Configuration

archon/agent_tools.py

@@ -12,35 +12,24 @@
 
 # Imports for hierarchical retrieval
 try:
-    from archon.llms_txt.retrieval.query_processor import QueryProcessor
-    # from archon.llms_txt.retrieval.ranking import HierarchicalRanker # Not needed for direct search
-    # from archon.llms_txt.retrieval.response_builder import ResponseBuilder # Not needed for direct search
-    # from archon.llms_txt.retrieval.retrieval_manager import RetrievalManager # Not needed for direct search
+    # from archon.llms_txt.retrieval.query_processor import QueryProcessor # Handled by EmbeddingManager now
+    # from archon.llms_txt.retrieval.ranking import HierarchicalRanker
+    # from archon.llms_txt.retrieval.response_builder import ResponseBuilder
+    # from archon.llms_txt.retrieval.retrieval_manager import RetrievalManager
     from archon.llms_txt.vector_db.supabase_manager import SupabaseManager
+    from archon.llms_txt.vector_db.embedding_manager import EmbeddingManager # Import EmbeddingManager
     HIERARCHICAL_IMPORTS_AVAILABLE = True
 except ImportError as e:
     print(f"Warning: Failed to import hierarchical retrieval components: {e}. Hierarchical retrieval disabled.")
     HIERARCHICAL_IMPORTS_AVAILABLE = False
     # Define dummy classes if imports fail to prevent NameErrors later
-    class QueryProcessor: pass
+    class QueryProcessor: pass # Keep dummy for now if needed elsewhere
     class SupabaseManager: pass
+    class EmbeddingManager: pass # Add dummy for EmbeddingManager
 
+# Remove the old get_embedding function and embedding_model global variable
 
-embedding_model = get_env_var('EMBEDDING_MODEL') or 'text-embedding-3-small'
-
-async def get_embedding(text: str, embedding_client: AsyncOpenAI) -> List[float]:
-    """Get embedding vector from OpenAI."""
-    try:
-        response = await embedding_client.embeddings.create(
-            model=embedding_model,
-            input=text
-        )
-        return response.data[0].embedding
-    except Exception as e:
-        print(f"Error getting embedding: {e}")
-        return [0] * 1536  # Return zero vector on error
-
-async def retrieve_relevant_documentation_tool(supabase: Client, embedding_client: AsyncOpenAI, user_query: str) -> str:
+async def retrieve_relevant_documentation_tool(supabase: Client, embedding_manager: EmbeddingManager, user_query: str) -> str: # Changed embedding_client to embedding_manager
     """
     Retrieve relevant documentation chunks based on the query.
     Conditionally uses either standard RAG from 'site_pages' or hierarchical RAG
@@ -58,14 +47,11 @@ async def retrieve_relevant_documentation_tool(supabase: Client, embedding_clien
             try:
                 # 1. Initialize Hierarchical Components
                 # SupabaseManager handles its own env loading and client creation
-                db_manager = SupabaseManager()
-                # QueryProcessor handles its own env loading and embedder creation
-                # We need to ensure QueryProcessor uses the *same* embedding client/model
-                # Let's reuse the existing get_embedding function for simplicity and consistency
-                # query_processor = QueryProcessor(embedding_client=embedding_client) # Pass existing client
+                db_manager = SupabaseManager() # Assumes SupabaseManager uses its own env loader correctly
 
-                # 2. Process Query to get Embedding
-                query_embedding = await get_embedding(user_query, embedding_client)
+                # 2. Process Query to get Embedding using EmbeddingManager
+                print(f"DEBUG: Generating embedding for query '{user_query}' using {embedding_manager.provider} provider.")
+                query_embedding = embedding_manager.generate_embedding(user_query) # Use manager's method
 
                 if not query_embedding or len(query_embedding) < 10: # Basic check
                     return "Error: Failed to generate embedding for hierarchical search."
@@ -108,8 +94,9 @@ async def retrieve_relevant_documentation_tool(supabase: Client, embedding_clien
 
         else: # Default to 'site_pages'
             print(f"--- Performing Standard Retrieval ('site_pages') for: '{user_query}' ---")
-            # Get the embedding for the query
-            query_embedding = await get_embedding(user_query, embedding_client)
+            # Get the embedding for the query using EmbeddingManager
+            print(f"DEBUG: Generating embedding for query '{user_query}' using {embedding_manager.provider} provider.")
+            query_embedding = embedding_manager.generate_embedding(user_query) # Use manager's method
 
             # Query Supabase for relevant documents using the existing RPC
             result = supabase.rpc(

archon/archon_graph.py

@@ -22,6 +22,7 @@
 # Add the parent directory to Python path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from archon.pydantic_ai_coder import pydantic_ai_coder, PydanticAIDeps
+from archon.llms_txt.vector_db.embedding_manager import EmbeddingManager # Import EmbeddingManager
 from archon.refiner_agents.prompt_refiner_agent import prompt_refiner_agent
 from archon.refiner_agents.tools_refiner_agent import tools_refiner_agent, ToolsRefinerDeps
 from archon.refiner_agents.agent_refiner_agent import agent_refiner_agent, AgentRefinerDeps
@@ -114,11 +115,15 @@ async def define_scope_with_reasoner(state: AgentState):
     return {"scope": scope}
 
 # Coding Node with Feedback Handling
-async def coder_agent(state: AgentState, writer):    
-    # Prepare dependencies
+async def coder_agent(state: AgentState, writer):
+    # Initialize EmbeddingManager (it loads its own env vars)
+    # Note: Consider initializing this once outside the node if possible for efficiency
+    embedding_manager = EmbeddingManager()
+
+    # Prepare dependencies, passing the embedding_manager
     deps = PydanticAIDeps(
         supabase=supabase,
-        embedding_client=embedding_client,
+        embedding_manager=embedding_manager, # Pass the initialized EmbeddingManager
         reasoner_output=state['scope']
     )
 

archon/llms_txt/retrieval/query_processor.py

@@ -1,5 +1,5 @@
 import re
-from archon.llms_txt.vector_db.embedding_manager import OpenAIEmbeddingGenerator
+from archon.llms_txt.vector_db.embedding_manager import EmbeddingManager
 from typing import Any, Dict, List, Optional, Tuple
 
 class QueryProcessor:
@@ -13,10 +13,10 @@ def __init__(self):
         Initializes the QueryProcessor and the embedding generator.
         """
         try:
-            self.embedder = OpenAIEmbeddingGenerator()
-            print("QueryProcessor: OpenAIEmbeddingGenerator initialized.")
+            self.embedder = EmbeddingManager()
+            print("QueryProcessor: EmbeddingManager initialized.")
         except Exception as e:
-            print(f"Error initializing OpenAIEmbeddingGenerator in QueryProcessor: {e}")
+            print(f"Error initializing EmbeddingManager in QueryProcessor: {e}")
             self.embedder = None # Handle initialization failure
 
     def extract_contextual_info(self, query: str) -> Dict[str, Any]: