Refactor context retrieval to include advanced model for improved query handling

Author: dcloud347

prometheus/lang_graph/nodes/context_retrieval_subgraph_node.py

@@ -13,7 +13,8 @@
 class ContextRetrievalSubgraphNode:
     def __init__(
         self,
-        model: BaseChatModel,
+        base_model: BaseChatModel,
+        advanced_model: BaseChatModel,
         kg: KnowledgeGraph,
         local_path: str,
         query_key_name: str,
@@ -22,7 +23,8 @@ def __init__(
     ):
         self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}")
         self.context_retrieval_subgraph = ContextRetrievalSubgraph(
-            model=model,
+            base_model=base_model,
+            advanced_model=advanced_model,
             kg=kg,
             local_path=local_path,
             repository_id=repository_id,

prometheus/lang_graph/subgraphs/bug_get_regression_tests_subgraph.py

@@ -55,6 +55,7 @@ def __init__(
         # Step 2: Retrieve relevant code/documentation context from the knowledge graph
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
             base_model,
+            advanced_model,
             kg,
             git_repo.playground_path,
             "select_regression_query",

prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py

@@ -63,6 +63,7 @@ def __init__(
         # Step 2: Retrieve relevant code/documentation context from the knowledge graph
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
             base_model,
+            advanced_model,
             kg,
             git_repo.playground_path,
             "bug_reproducing_query",

prometheus/lang_graph/subgraphs/context_retrieval_subgraph.py

@@ -68,7 +68,8 @@ class ContextRetrievalSubgraph:
 
     def __init__(
         self,
-        model: BaseChatModel,
+        base_model: BaseChatModel,
+        advanced_model: BaseChatModel,
         kg: KnowledgeGraph,
         local_path: str,
         repository_id: int,
@@ -77,19 +78,20 @@ def __init__(
         Initializes the context retrieval subgraph.
 
         Args:
-            model (BaseChatModel): The LLM used for context selection and refinement.
+            base_model (BaseChatModel): The LLM used for context selection and refinement.
+            advanced_model (BaseChatModel): The LLM used for advanced tasks like query refinement and context extraction.
             kg (KnowledgeGraph): Knowledge graph instance
             local_path (str): Local path to the codebase for context extraction.
             repository_id (int): Repository ID for memory storage
         """
         # Step 1: Refine query into structured format
-        context_refine_node = ContextRefineNode(model=model, kg=kg)
+        context_refine_node = ContextRefineNode(model=advanced_model, kg=kg)
 
         # Step 2: Retrieve contexts from semantic memory (Athena)
         memory_retrieval_node = MemoryRetrievalNode(repository_id=repository_id)
 
         # Step 3: Extract relevant contexts from explored_context
-        context_extraction_node = ContextExtractionNode(model=model, root_path=local_path)
+        context_extraction_node = ContextExtractionNode(model=advanced_model, root_path=local_path)
 
         # Step 4: Store new contexts to memory
         memory_storage_node = MemoryStorageNode(repository_id=repository_id)
@@ -101,7 +103,7 @@ def __init__(
         add_context_refined_query_message_node = AddContextRefinedQueryMessageNode()
 
         # Step 7: Query knowledge graph (Neo4j) using LLM tools
-        context_provider_node = ContextProviderNode(model=model, kg=kg, local_path=local_path)
+        context_provider_node = ContextProviderNode(model=base_model, kg=kg, local_path=local_path)
         context_provider_tools = ToolNode(
             tools=context_provider_node.tools,
             name="context_provider_tools",

prometheus/lang_graph/subgraphs/issue_classification_subgraph.py

@@ -22,7 +22,8 @@ def __init__(
     ):
         issue_classification_context_message_node = IssueClassificationContextMessageNode()
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
-            model=model,
+            base_model=model,
+            advanced_model=model,
             kg=kg,
             local_path=local_path,
             query_key_name="issue_classification_query",

prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py

@@ -37,7 +37,8 @@ def __init__(
     ):
         issue_bug_context_message_node = IssueBugContextMessageNode()
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
-            model=base_model,
+            base_model=base_model,
+            advanced_model=advanced_model,
             kg=kg,
             local_path=git_repo.playground_path,
             query_key_name="bug_fix_query",

prometheus/lang_graph/subgraphs/issue_question_subgraph.py

@@ -37,7 +37,8 @@ def __init__(
         # Step 1: Retrieve relevant context based on the issue details
         issue_question_context_message_node = IssueQuestionContextMessageNode()
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
-            model=base_model,
+            base_model=base_model,
+            advanced_model=advanced_model,
             kg=kg,
             local_path=git_repo.playground_path,
             query_key_name="question_query",

prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py

@@ -73,7 +73,8 @@ def __init__(
         # Phase 1: Retrieve context related to the bug
         issue_bug_context_message_node = IssueBugContextMessageNode()
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
-            model=base_model,
+            base_model=base_model,
+            advanced_model=advanced_model,
             kg=kg,
             local_path=git_repo.playground_path,
             query_key_name="bug_fix_query",