Feat/llm provider query (#1735)

* Add ModelProvider to Query package.

* Spellcheck + others

* Semver

* Fix tests

* Format

* Fix Pyright

* Fix tests

* Fix for smoke tests

Author: AlonsoGuevara

.semversioner/next-release/patch-20250224223906608414.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Use ModelProvider for query module"
+}

dictionary.txt

@@ -120,17 +120,16 @@ unhot
 groupby
 retryer
 agenerate
-aembed
-dedupe
 dropna
-dtypes
 notna
 
 # LLM Terms
 AOAI
 embedder
 llm
 llms
+achat
+aembed
 
 # Galaxy-Brain Terms
 Unipartite

graphrag/callbacks/query_callbacks.py

@@ -28,3 +28,6 @@ def on_reduce_response_start(
 
     def on_reduce_response_end(self, reduce_response_output: str) -> None:
         """Handle the end of reduce operation."""
+
+    def on_llm_new_token(self, token) -> None:
+        """Handle when a new token is generated."""

graphrag/config/defaults.py

@@ -48,6 +48,8 @@ class BasicSearchDefaults:
     n: int = 1
     max_tokens: int = 12_000
     llm_max_tokens: int = 2000
+    chat_model_id: str = DEFAULT_CHAT_MODEL_ID
+    embedding_model_id: str = DEFAULT_EMBEDDING_MODEL_ID
 
 
 @dataclass
@@ -122,7 +124,9 @@ class DriftSearchDefaults:
     local_search_temperature: float = 0
     local_search_top_p: float = 1
     local_search_n: int = 1
-    local_search_llm_max_gen_tokens: int = 12_000
+    local_search_llm_max_gen_tokens: int = 4_096
+    chat_model_id: str = DEFAULT_CHAT_MODEL_ID
+    embedding_model_id: str = DEFAULT_EMBEDDING_MODEL_ID
 
 
 @dataclass
@@ -239,6 +243,8 @@ class GlobalSearchDefaults:
     dynamic_search_use_summary: bool = False
     dynamic_search_concurrent_coroutines: int = 16
     dynamic_search_max_level: int = 2
+    chat_model_id: str = DEFAULT_CHAT_MODEL_ID
+    embedding_model_id: str = DEFAULT_EMBEDDING_MODEL_ID
 
 
 @dataclass
@@ -305,6 +311,8 @@ class LocalSearchDefaults:
     n: int = 1
     max_tokens: int = 12_000
     llm_max_tokens: int = 2000
+    chat_model_id: str = DEFAULT_CHAT_MODEL_ID
+    embedding_model_id: str = DEFAULT_EMBEDDING_MODEL_ID
 
 
 @dataclass

graphrag/config/init_content.py

@@ -145,18 +145,26 @@
 ## See the config docs: https://microsoft.github.io/graphrag/config/yaml/#query
 
 local_search:
+  chat_model_id: {graphrag_config_defaults.local_search.chat_model_id}
+  embedding_model_id: {graphrag_config_defaults.local_search.embedding_model_id}
   prompt: "prompts/local_search_system_prompt.txt"
 
 global_search:
+  chat_model_id: {graphrag_config_defaults.global_search.chat_model_id}
+  embedding_model_id: {graphrag_config_defaults.global_search.embedding_model_id}
   map_prompt: "prompts/global_search_map_system_prompt.txt"
   reduce_prompt: "prompts/global_search_reduce_system_prompt.txt"
   knowledge_prompt: "prompts/global_search_knowledge_system_prompt.txt"
 
 drift_search:
+  chat_model_id: {graphrag_config_defaults.drift_search.chat_model_id}
+  embedding_model_id: {graphrag_config_defaults.drift_search.embedding_model_id}
   prompt: "prompts/drift_search_system_prompt.txt"
   reduce_prompt: "prompts/drift_search_reduce_prompt.txt"
 
 basic_search:
+  chat_model_id: {graphrag_config_defaults.basic_search.chat_model_id}
+  embedding_model_id: {graphrag_config_defaults.basic_search.embedding_model_id}
   prompt: "prompts/basic_search_system_prompt.txt"
 """
 

graphrag/config/models/basic_search_config.py

@@ -15,6 +15,14 @@ class BasicSearchConfig(BaseModel):
         description="The basic search prompt to use.",
         default=graphrag_config_defaults.basic_search.prompt,
     )
+    chat_model_id: str = Field(
+        description="The model ID to use for basic search.",
+        default=graphrag_config_defaults.basic_search.chat_model_id,
+    )
+    embedding_model_id: str = Field(
+        description="The model ID to use for text embeddings.",
+        default=graphrag_config_defaults.basic_search.embedding_model_id,
+    )
     text_unit_prop: float = Field(
         description="The text unit proportion.",
         default=graphrag_config_defaults.basic_search.text_unit_prop,

graphrag/config/models/drift_search_config.py

@@ -19,6 +19,14 @@ class DRIFTSearchConfig(BaseModel):
         description="The drift search reduce prompt to use.",
         default=graphrag_config_defaults.drift_search.reduce_prompt,
     )
+    chat_model_id: str = Field(
+        description="The model ID to use for drift search.",
+        default=graphrag_config_defaults.drift_search.chat_model_id,
+    )
+    embedding_model_id: str = Field(
+        description="The model ID to use for drift search.",
+        default=graphrag_config_defaults.drift_search.embedding_model_id,
+    )
     temperature: float = Field(
         description="The temperature to use for token generation.",
         default=graphrag_config_defaults.drift_search.temperature,

graphrag/config/models/global_search_config.py

@@ -19,6 +19,10 @@ class GlobalSearchConfig(BaseModel):
         description="The global search reducer to use.",
         default=graphrag_config_defaults.global_search.reduce_prompt,
     )
+    chat_model_id: str = Field(
+        description="The model ID to use for global search.",
+        default=graphrag_config_defaults.global_search.chat_model_id,
+    )
     knowledge_prompt: str | None = Field(
         description="The global search general prompt to use.",
         default=graphrag_config_defaults.global_search.knowledge_prompt,

graphrag/config/models/local_search_config.py

@@ -15,6 +15,14 @@ class LocalSearchConfig(BaseModel):
         description="The local search prompt to use.",
         default=graphrag_config_defaults.local_search.prompt,
     )
+    chat_model_id: str = Field(
+        description="The model ID to use for local search.",
+        default=graphrag_config_defaults.local_search.chat_model_id,
+    )
+    embedding_model_id: str = Field(
+        description="The model ID to use for text embeddings.",
+        default=graphrag_config_defaults.local_search.embedding_model_id,
+    )
     text_unit_prop: float = Field(
         description="The text unit proportion.",
         default=graphrag_config_defaults.local_search.text_unit_prop,

graphrag/index/operations/embed_text/strategies/openai.py

@@ -88,7 +88,7 @@ async def _execute(
 ) -> list[list[float]]:
     async def embed(chunk: list[str]):
         async with semaphore:
-            chunk_embeddings = await model.embed(chunk)
+            chunk_embeddings = await model.aembed_batch(chunk)
             result = np.array(chunk_embeddings)
             tick(1)
         return result