Update kig_core/graph_operations.py

kig_core/graph_operations.py

@@ -10,7 +10,7 @@ from langchain_core.pydantic_v1 import Field, BaseModel as V1BaseModel # For gra
 
 from .config import settings
 from .graph_client import neo4j_client # Use the central client
-from .llm_interface import get_llm
+from .llm_interface import get_llm, invoke_llm
 from .prompts import (
     CYPHER_GENERATION_PROMPT, CONCEPT_SELECTION_PROMPT,
     BINARY_GRADER_PROMPT, SCORE_GRADER_PROMPT
@@ -47,7 +47,7 @@ def generate_cypher_auto(question: str) -> str:
         | StrOutputParser()
         | extract_cypher
     )
-    return chain.invoke(question)
+    return llm_invoke(chain,question)
 
 def generate_cypher_guided(question: str, plan_step: int) -> str:
     """Generates Cypher using the 'guided' method based on concepts."""
@@ -64,7 +64,7 @@ def generate_cypher_guided(question: str, plan_step: int) -> str:
             | concept_llm
             | StrOutputParser()
         )
-        selected_concept = concept_chain.invoke({
+        selected_concept = llm_invoke(concept_chain,{
             "question": question,
             "concepts": "\n".join(concepts)
         }).strip()
@@ -179,7 +179,7 @@ def evaluate_documents(
             formatted_doc = format_doc_for_llm(doc)
             if not formatted_doc.strip(): continue
             try:
-                result = binary_grader.invoke({"question": query, "document": formatted_doc})
+                result = llm_invoke(binary_grader,{"question": query, "document": formatted_doc})
                 logger.debug(f"Binary grader result for doc '{doc.get('title', 'N/A')}': {result}")
                 if result and 'yes' in result.lower():
                     valid_docs_with_scores.append((doc, 1.0)) # Score 1.0 for relevant
@@ -193,7 +193,7 @@ def evaluate_documents(
             formatted_doc = format_doc_for_llm(doc)
             if not formatted_doc.strip(): continue
             try:
-                result: GradeDocumentsScore = score_grader.invoke({"query": query, "document": formatted_doc})
+                result: GradeDocumentsScore = llm_invoke(score_grader,{"query": query, "document": formatted_doc})
                 logger.debug(f"Score grader result for doc '{doc.get('title', 'N/A')}': Score={result.score}, Rationale={result.rationale}")
                 if result.score >= settings.eval_threshold:
                     valid_docs_with_scores.append((doc, result.score))