Add close/aclose methods to LLMInterfaceV2

AGENTS.md

@@ -28,7 +28,7 @@ BUNDLING: Hatchling (wheel builder)
 TESTING FRAMEWORKS: pytest, coverage
 
 TEST COMMANDS:
-- Unit tests: `poetry run pytest tests/unit/` (per project instructions) or `uv run pytest tests/unit`
+- Unit tests: `uv run pytest tests/unit/` (per project instructions) or `uv run pytest tests/unit`
 - E2E tests: `uv run pytest tests/e2e`
 - Coverage check: `uv run coverage run -m pytest tests/unit && uv run coverage report --fail-under=90`
 

examples/build_graph/automatic_schema_extraction/simple_kg_builder_schema_from_pdf.py

@@ -63,7 +63,7 @@ async def run_kg_pipeline_with_auto_schema() -> None:
 
     finally:
         # Close connections
-        await llm.async_client.close()
+        await llm.aclose()
         driver.close()
 
 

examples/build_graph/automatic_schema_extraction/simple_kg_builder_schema_from_text.py

@@ -81,7 +81,7 @@ async def run_kg_pipeline_with_auto_schema() -> None:
 
     finally:
         # Close connections
-        await llm.async_client.close()
+        await llm.aclose()
         driver.close()
 
 

examples/build_graph/simple_kg_builder_from_pdf.py

@@ -63,12 +63,9 @@ async def define_and_run_pipeline(
 
 
 async def main() -> PipelineResult:
-    llm = OpenAILLM(
-        model_name="gpt-5",
-    )
     with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
-        res = await define_and_run_pipeline(driver, llm)
-    await llm.async_client.close()
+        async with OpenAILLM(model_name="gpt-5") as llm:
+            res = await define_and_run_pipeline(driver, llm)
     return res
 
 

examples/build_graph/simple_kg_builder_from_text.py

@@ -94,12 +94,10 @@ async def define_and_run_pipeline(
 
 
 async def main() -> PipelineResult:
-    llm = OpenAILLM(
-        model_name="gpt-5",
-    )
     with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
-        res = await define_and_run_pipeline(driver, llm)
-    await llm.async_client.close()
+        async with OpenAILLM(model_name="gpt-5") as llm:
+            res = await define_and_run_pipeline(driver, llm)
+    await llm.aclose()
     return res
 
 

examples/customize/build_graph/components/schema_builders/schema_from_text.py

@@ -61,7 +61,7 @@ async def extract_and_save_schema() -> None:
 
     # Define LLM parameters
     llm_model_params = {
-        "max_tokens": 2000,
+        "max_completion_tokens": 2000,
         "response_format": {"type": "json_object"},
         "temperature": 0,  # Lower temperature for more consistent output
     }
@@ -104,7 +104,7 @@ async def extract_and_save_schema() -> None:
 
     finally:
         # Close the LLM client
-        await llm.async_client.close()
+        await llm.aclose()
 
 
 async def main() -> None:

examples/customize/build_graph/pipeline/kg_builder_from_pdf.py

@@ -92,6 +92,7 @@ async def define_and_run_pipeline(
         LLMEntityRelationExtractor(
             llm=llm,
             on_error=OnError.RAISE,
+            use_structured_output=True,
         ),
         "extractor",
     )
@@ -126,19 +127,18 @@ async def define_and_run_pipeline(
 
 
 async def main() -> PipelineResult:
-    llm = OpenAILLM(
-        model_name="gpt-5",
-        model_params={
-            "max_tokens": 2000,
-            "response_format": {"type": "json_object"},
-        },
-    )
-    driver = neo4j.GraphDatabase.driver(
-        "bolt://localhost:7687", auth=("neo4j", "password")
-    )
-    res = await define_and_run_pipeline(driver, llm)
-    driver.close()
-    await llm.async_client.close()
+    res = None
+    try:
+        llm = OpenAILLM(
+            model_name="gpt-5",
+        )
+        driver = neo4j.GraphDatabase.driver(
+            "bolt://localhost:7687", auth=("neo4j", "password")
+        )
+        res = await define_and_run_pipeline(driver, llm)
+    finally:
+        driver.close()
+        await llm.aclose()
     return res
 
 

examples/customize/build_graph/pipeline/kg_builder_from_text.py

@@ -68,6 +68,7 @@ async def define_and_run_pipeline(
         LLMEntityRelationExtractor(
             llm=llm,
             on_error=OnError.RAISE,
+            use_structured_output=True,
         ),
         "extractor",
     )
@@ -142,19 +143,18 @@ async def define_and_run_pipeline(
 
 
 async def main() -> PipelineResult:
-    llm = OpenAILLM(
-        model_name="gpt-5",
-        model_params={
-            "max_tokens": 1000,
-            "response_format": {"type": "json_object"},
-        },
-    )
-    driver = neo4j.GraphDatabase.driver(
-        "bolt://localhost:7687", auth=("neo4j", "password")
-    )
-    res = await define_and_run_pipeline(driver, llm)
-    driver.close()
-    await llm.async_client.close()
+    res = None
+    try:
+        llm = OpenAILLM(
+            model_name="gpt-5",
+        )
+        driver = neo4j.GraphDatabase.driver(
+            "bolt://localhost:7687", auth=("neo4j", "password")
+        )
+        res = await define_and_run_pipeline(driver, llm)
+    finally:
+        driver.close()
+        await llm.aclose()
     return res
 
 

examples/customize/build_graph/pipeline/kg_builder_two_documents_entity_resolution.py

@@ -66,6 +66,7 @@ async def define_and_run_pipeline(
         LLMEntityRelationExtractor(
             llm=llm,
             on_error=OnError.IGNORE,
+            use_structured_output=True,
         ),
         "extractor",
     )
@@ -140,16 +141,15 @@ async def main() -> None:
     llm = OpenAILLM(
         model_name="gpt-5",
         model_params={
-            "max_tokens": 1000,
-            "response_format": {"type": "json_object"},
+            "max_completion_tokens": 1000,
         },
     )
     driver = neo4j.GraphDatabase.driver(
         "bolt://localhost:7687", auth=("neo4j", "password")
     )
     await define_and_run_pipeline(driver, llm)
     driver.close()
-    await llm.async_client.close()
+    await llm.aclose()
 
 
 if __name__ == "__main__":

examples/customize/build_graph/pipeline/text_to_lexical_graph_to_entity_graph_single_pipeline.py

@@ -185,7 +185,7 @@ async def main(driver: neo4j.Driver) -> PipelineResult:
         lexical_graph_config,
         text,
     )
-    await llm.async_client.close()
+    await llm.aclose()
     return res
 
 

examples/customize/build_graph/pipeline/text_to_lexical_graph_to_entity_graph_two_pipelines.py

@@ -204,7 +204,7 @@ async def main(driver: neo4j.Driver) -> PipelineResult:
     res = await read_chunk_and_perform_entity_extraction(
         driver, llm, lexical_graph_config
     )
-    await llm.async_client.close()
+    await llm.aclose()
     return res
 
 

examples/customize/llms/anthropic_llm.py

@@ -3,10 +3,10 @@
 # set api key here on in the ANTHROPIC_API_KEY env var
 api_key = None
 
-llm = AnthropicLLM(
+with AnthropicLLM(
     model_name="claude-3-opus-20240229",
     model_params={"max_tokens": 1000},  # max_tokens must be specified
     api_key=api_key,
-)
-res: LLMResponse = llm.invoke("say something")
-print(res.content)
+) as llm:
+    res: LLMResponse = llm.invoke("say something")
+    print(res.content)

examples/customize/llms/cohere_llm.py

@@ -3,9 +3,9 @@
 # set api key here on in the CO_API_KEY env var
 api_key = None
 
-llm = CohereLLM(
+with CohereLLM(
     model_name="command-r",
     api_key=api_key,
-)
-res: LLMResponse = llm.invoke("say something")
-print(res.content)
+) as llm:
+    res: LLMResponse = llm.invoke("say something")
+    print(res.content)

examples/customize/llms/llm_with_message_history.py

@@ -10,33 +10,22 @@
 # set api key here on in the OPENAI_API_KEY env var
 api_key = None
 
-llm = OpenAILLM(model_name="gpt-5", api_key=api_key)
-
 questions = [
     "What are some movies Tom Hanks starred in?",
     "Is he also a director?",
     "Wow, that's impressive. And what about his personal life, does he have children?",
 ]
 
 history: list[dict[str, str]] = []
-for question in questions:
-    res: LLMResponse = llm.invoke(
-        question,
-        message_history=history,  # type: ignore
-    )
-    history.append(
-        {
-            "role": "user",
-            "content": question,
-        }
-    )
-    history.append(
-        {
-            "role": "assistant",
-            "content": res.content,
-        }
-    )
+with OpenAILLM(model_name="gpt-5", api_key=api_key) as llm:
+    for question in questions:
+        res: LLMResponse = llm.invoke(
+            question,
+            message_history=history,  # type: ignore
+        )
+        history.append({"role": "user", "content": question})
+        history.append({"role": "assistant", "content": res.content})
 
-    print("#" * 50, question)
-    print(res.content)
-    print("#" * 50)
+        print("#" * 50, question)
+        print(res.content)
+        print("#" * 50)

examples/customize/llms/llm_with_neo4j_message_history.py

@@ -36,24 +36,15 @@
 
 history = Neo4jMessageHistory(session_id="123", driver=driver, window=10)
 
-for question in questions:
-    res: LLMResponse = llm.invoke(
-        question,
-        message_history=history,
-    )
-    history.add_message(
-        {
-            "role": "user",
-            "content": question,
-        }
-    )
-    history.add_message(
-        {
-            "role": "assistant",
-            "content": res.content,
-        }
-    )
-
-    print("#" * 50, question)
-    print(res.content)
-    print("#" * 50)
+with OpenAILLM(model_name="gpt-5", api_key=api_key) as llm:
+    for question in questions:
+        res: LLMResponse = llm.invoke(
+            question,
+            message_history=history,
+        )
+        history.add_message({"role": "user", "content": question})
+        history.add_message({"role": "assistant", "content": res.content})
+
+        print("#" * 50, question)
+        print(res.content)
+        print("#" * 50)

examples/customize/llms/llm_with_system_instructions.py

@@ -8,15 +8,11 @@
 # set api key here on in the OPENAI_API_KEY env var
 api_key = None
 
-llm = OpenAILLM(
-    model_name="gpt-5",
-    api_key=api_key,
-)
-
 question = "How fast is Santa Claus during the Christmas eve?"
 
-res: LLMResponse = llm.invoke(
-    question,
-    system_instruction="Answer with a serious tone",
-)
-print(res.content)
+with OpenAILLM(model_name="gpt-5", api_key=api_key) as llm:
+    res: LLMResponse = llm.invoke(
+        question,
+        system_instruction="Answer with a serious tone",
+    )
+    print(res.content)

examples/customize/llms/mistalai_llm.py

@@ -3,8 +3,8 @@
 # set api key here on in the MISTRAL_API_KEY env var
 api_key = None
 
-llm = MistralAILLM(
+with MistralAILLM(
     model_name="mistral-small-latest",
     api_key=api_key,
-)
-llm.invoke("say something")
+) as llm:
+    llm.invoke("say something")

examples/customize/llms/ollama_llm.py

@@ -4,10 +4,10 @@
 
 from neo4j_graphrag.llm import LLMResponse, OllamaLLM
 
-llm = OllamaLLM(
+with OllamaLLM(
     model_name="<model_name>",
     # model_params={"options": {"temperature": 0}, "format": "json"},
     # host="...",  # if using a remote server
-)
-res: LLMResponse = llm.invoke("What is the additive color model?")
-print(res.content)
+) as llm:
+    res: LLMResponse = llm.invoke("What is the additive color model?")
+    print(res.content)