Change GraphRAG search parameter query to query_text (#89)

* Fix doc for GraphRAG (arg is called query and not query_text=

* Rename GraphRAG search parameter query to query_text for consistency with Retriever interface - and let the possibility to add query_vector param later on if requested

* Update CHANGELOG

* Check disk space before install

* Check container size

* Check docker image size

* Test with one single neo/python version

Author: stellasia

CHANGELOG.md

@@ -8,6 +8,9 @@
 -   Corrected initialization to allow specifying the embedding model name.
 -   Removed sentence_transformers from embeddings/__init__.py to avoid ImportError when the package is not installed.
 
+### Changed
+- `GraphRAG.search` method first parameter has been renamed `query_text` (was `query`) for consistency with the retrievers interface.
+
 ## 0.3.0
 
 ### Added

README.md

@@ -128,7 +128,7 @@ rag = GraphRAG(retriever=retriever, llm=llm)
 
 # Query the graph
 query_text = "How do I do similarity search in Neo4j?"
-response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
+response = rag.search(query=query_text, retriever_config={"top_k": 5})
 print(response.answer)
 ```
 

examples/graphrag_custom_prompt.py

@@ -60,7 +60,7 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:
     {context}
 
     Question:
-    {query}
+    {query_text}
 
     Answer:
     """

src/neo4j_genai/generation/graphrag.py

@@ -53,7 +53,7 @@ def __init__(
 
     def search(
         self,
-        query: str,
+        query_text: str,
         examples: str = "",
         retriever_config: Optional[dict[str, Any]] = None,
         return_context: bool = False,
@@ -64,7 +64,7 @@ def search(
         3. Generation: answer generation with LLM
 
         Args:
-            query (str): The user question
+            query_text (str): The user question
             examples: Examples added to the LLM prompt.
             retriever_config (Optional[dict]): Parameters passed to the retriever
                 search method; e.g.: top_k
@@ -76,20 +76,20 @@ def search(
         """
         try:
             validated_data = RagSearchModel(
-                query=query,
+                query_text=query_text,
                 examples=examples,
                 retriever_config=retriever_config or {},
                 return_context=return_context,
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors())
-        query = validated_data.query
+        query_text = validated_data.query_text
         retriever_result: RetrieverResult = self.retriever.search(
-            query_text=query, **validated_data.retriever_config
+            query_text=query_text, **validated_data.retriever_config
         )
         context = "\n".join(item.content for item in retriever_result.items)
         prompt = self.prompt_template.format(
-            query=query, context=context, examples=validated_data.examples
+            query_text=query_text, context=context, examples=validated_data.examples
         )
         logger.debug(f"RAG: retriever_result={retriever_result}")
         logger.debug(f"RAG: prompt={prompt}")

src/neo4j_genai/generation/prompts.py

@@ -87,14 +87,14 @@ class RagTemplate(PromptTemplate):
 {examples}
 
 Question:
-{query}
+{query_text}
 
 Answer:
 """
-    EXPECTED_INPUTS = ["context", "query", "examples"]
+    EXPECTED_INPUTS = ["context", "query_text", "examples"]
 
-    def format(self, query: str, context: str, examples: str) -> str:
-        return super().format(query=query, context=context, examples=examples)
+    def format(self, query_text: str, context: str, examples: str) -> str:
+        return super().format(query_text=query_text, context=context, examples=examples)
 
 
 class Text2CypherTemplate(PromptTemplate):

src/neo4j_genai/generation/types.py

@@ -39,7 +39,7 @@ def check_llm(cls, value: Any) -> Any:
 
 
 class RagSearchModel(BaseModel):
-    query: str
+    query_text: str
     examples: str = ""
     retriever_config: dict[str, Any] = {}
     return_context: bool = False

tests/e2e/test_graphrag_e2e.py

@@ -52,7 +52,7 @@ def test_graphrag_happy_path(
     llm.invoke.return_value = LLMResponse(content="some text")
 
     result = rag.search(
-        query="biology",
+        query_text="biology",
         retriever_config={
             "top_k": 2,
         },
@@ -96,7 +96,7 @@ def test_graphrag_happy_path_return_context(
     llm.invoke.return_value = LLMResponse(content="some text")
 
     result = rag.search(
-        query="biology",
+        query_text="biology",
         retriever_config={
             "top_k": 2,
         },
@@ -142,7 +142,7 @@ def test_graphrag_happy_path_examples(
     llm.invoke.return_value = LLMResponse(content="some text")
 
     result = rag.search(
-        query="biology",
+        query_text="biology",
         retriever_config={
             "top_k": 2,
         },
@@ -186,7 +186,7 @@ def test_graphrag_llm_error(
 
     with pytest.raises(LLMGenerationError):
         rag.search(
-            query="biology",
+            query_text="biology",
         )
 
 
@@ -203,5 +203,5 @@ def test_graphrag_retrieval_error(
 
     with pytest.raises(TypeError):
         rag.search(
-            query="biology",
+            query_text="biology",
         )

tests/unit/test_graphrag.py

@@ -25,7 +25,9 @@
 
 def test_graphrag_prompt_template() -> None:
     template = RagTemplate()
-    prompt = template.format(context="my context", query="user's query", examples="")
+    prompt = template.format(
+        context="my context", query_text="user's query", examples=""
+    )
     assert (
         prompt
         == """Answer the user question using the following context