fix: stop demo query streams from dying mid-flight (research#86)

.env.example

@@ -75,6 +75,19 @@ FALKORDB_URL=redis://localhost:6379/0  # REQUIRED - change to your FalkorDB URL
 #   COMPLETION_MODEL=openai/gpt-4.1
 #   EMBEDDING_MODEL=openai/text-embedding-ada-002
 
+# Wall-clock ceiling for a single agent LLM call, in seconds (default 90).
+# Passed to litellm, which aborts the HTTP request — a hung provider then
+# surfaces as an error instead of stalling the response stream.
+# LLM_TIMEOUT=90
+#
+# Calls slower than this are logged at WARNING (default 20).
+# LLM_SLOW_CALL_THRESHOLD=20
+#
+# Retry budget per LLM call (default 1). LLM_TIMEOUT applies per attempt, so
+# this is pinned rather than left to the provider SDK and litellm defaults,
+# which each retry and together multiply the effective ceiling.
+# LLM_MAX_RETRIES=1
+
 # OpenAI - uses openai/gpt-4.1 and openai/text-embedding-ada-002
 # OPENAI_API_KEY=your_openai_api_key
 
@@ -100,7 +113,9 @@ FALKORDB_URL=redis://localhost:6379/0  # REQUIRED - change to your FalkorDB URL
 # Azure OpenAI (default fallback) - uses azure/gpt-4.1 and azure/text-embedding-ada-002
 # AZURE_API_KEY=your_azure_api_key
 # AZURE_API_BASE=https://your-resource.openai.azure.com/
-# AZURE_API_VERSION=2023-05-15
+# Must be 2025-03-01-preview or later — Graphiti memory writes use the
+# Azure Responses API, which rejects older api-versions with HTTP 400.
+# AZURE_API_VERSION=2025-03-01-preview
 
 # -----------------------------
 # OAuth configuration (optional — uncomment to enable login flows)

README.md

@@ -502,7 +502,7 @@ docker run -p 5000:5000 -it \
   -e FASTAPI_SECRET_KEY=your_secret_key \
   -e AZURE_API_KEY=your_azure_api_key \
   -e AZURE_API_BASE=https://your-resource.openai.azure.com/ \
-  -e AZURE_API_VERSION=2024-12-01-preview \
+  -e AZURE_API_VERSION=2025-03-01-preview \
   falkordb/queryweaver
 ```
 

api/agents/analysis_agent.py

@@ -38,7 +38,8 @@ def get_analysis(  # pylint: disable=too-many-arguments, too-many-positional-arg
         self.messages.append({"role": "user", "content": prompt})
 
         response = run_completion(
-            self.messages, self.custom_model, self.custom_api_key, temperature=0
+            self.messages, self.custom_model, self.custom_api_key,
+            label="analysis", temperature=0,
         )
         analysis = parse_response(response)
         if isinstance(analysis["ambiguities"], list):

api/agents/follow_up_agent.py

@@ -70,7 +70,8 @@ def generate_follow_up_question(
         try:
             response = run_completion(
                 [{"role": "user", "content": prompt}],
-                self.custom_model, self.custom_api_key, temperature=0.9
+                self.custom_model, self.custom_api_key,
+                label="followup", temperature=0.9,
             )
             return response.strip()
 

api/agents/healer_agent.py

@@ -10,9 +10,7 @@
 
 import re
 from typing import Dict, Callable, Any
-from litellm import completion
-from api.config import Config
-from .utils import parse_response
+from .utils import parse_response, run_completion
 
 
 class HealerAgent:
@@ -224,14 +222,12 @@ def heal_and_execute(  # pylint: disable=too-many-locals
 
         for attempt in range(self.max_healing_attempts):
             # Call LLM
-            response = completion(
-                model=Config.COMPLETION_MODEL,
-                messages=self.messages,
+            content = run_completion(
+                self.messages,
+                label=f"healer.attempt{attempt + 1}",
                 temperature=0.1,
                 max_tokens=2000
             )
-
-            content = response.choices[0].message.content
             self.messages.append({"role": "assistant", "content": content})
 
             # Parse response

api/agents/relevancy_agent.py

@@ -1,5 +1,6 @@
 """Relevancy agent for determining relevancy of queries to database schema."""
 
+import asyncio
 import json
 from .utils import BaseAgent, parse_response, run_completion
 
@@ -82,8 +83,14 @@ async def get_answer(self, user_question: str, database_desc: dict) -> dict:
             }
         )
 
-        answer = run_completion(
-            self.messages, self.custom_model, self.custom_api_key, temperature=0
+        # ``run_completion`` is synchronous. Awaiting it off-loop matters even
+        # though this method is already ``async``: the caller runs it as a task
+        # alongside table-finding, and a blocking call here would stall that
+        # task — and every other request — rather than overlap with it.
+        answer = await asyncio.to_thread(
+            run_completion,
+            self.messages, self.custom_model, self.custom_api_key,
+            label="relevancy", temperature=0,
         )
         self.messages.append({"role": "assistant", "content": answer})
         return parse_response(answer)

api/agents/response_formatter_agent.py

@@ -77,6 +77,7 @@ def format_response(self, user_query: str, sql_query: str,
 
         response = run_completion(
             messages, self.custom_model, self.custom_api_key,
+            label="formatter",
             temperature=0.3  # Slightly higher temperature for more natural responses
         )
         return response.strip()

api/agents/utils.py

@@ -1,29 +1,63 @@
 """Utility functions for agents."""
 
 import json
+import logging
+import time
 from typing import Any, Dict, List
 
 from litellm import completion
 from api.config import Config
 
 
-def run_completion(messages: List[Dict[str, str]], custom_model: str = None,
-                   custom_api_key: str = None, **kwargs) -> str:
+def run_completion(messages: List[Dict[str, str]], custom_model: str | None = None,
+                   custom_api_key: str | None = None, *, label: str = "llm",
+                   **kwargs) -> str:
     """Run an LLM completion with optional custom model/key overrides.
 
+    Applies ``Config.LLM_TIMEOUT`` per attempt and a pinned retry budget
+    unless the caller overrides them, and logs the call duration. Both exist because the 2026-07-29
+    demo failure was an LLM call that stalled with no timeout and left no
+    trace of how long it ran. ``label`` names the caller in those log lines
+    and is not forwarded to the provider.
+
     Returns the content string from the first choice.
     """
     completion_args = {
         "model": custom_model if custom_model else Config.COMPLETION_MODEL,
         "messages": messages,
         "top_p": 1,
+        "timeout": Config.LLM_TIMEOUT,
+        # ``timeout`` is per attempt, so the retry budget has to be pinned too
+        # or the effective ceiling becomes a multiple of it. litellm's outer
+        # retry loop is disabled in favour of the SDK-level count.
+        "max_retries": Config.LLM_MAX_RETRIES,
+        "num_retries": 0,
         **kwargs,
     }
 
     if custom_api_key:
         completion_args["api_key"] = custom_api_key
 
-    result = completion(**completion_args)
+    started = time.monotonic()
+    try:
+        result = completion(**completion_args)
+    except Exception:
+        logging.warning(
+            "llm_call label=%s model=%s duration=%.2fs outcome=error",
+            label, completion_args["model"], time.monotonic() - started,
+        )
+        raise
+    elapsed = time.monotonic() - started
+    logging.info(
+        "llm_call label=%s model=%s duration=%.2fs outcome=ok",
+        label, completion_args["model"], elapsed,
+    )
+    if elapsed >= Config.LLM_SLOW_CALL_THRESHOLD:
+        logging.warning(
+            "llm_call label=%s model=%s duration=%.2fs exceeded slow-call "
+            "threshold of %.0fs", label, completion_args["model"], elapsed,
+            Config.LLM_SLOW_CALL_THRESHOLD,
+        )
     return result.choices[0].message.content
 
 

api/config.py

@@ -133,6 +133,27 @@ class Config:
         COMPLETION_MODEL = _user_completion or "azure/gpt-4.1"
         EMBEDDING_MODEL_NAME = _user_embedding or "azure/text-embedding-ada-002"
 
+    # Wall-clock ceiling for a single agent LLM call, in seconds. Passed
+    # through to litellm, which aborts the underlying HTTP request — so a
+    # hung provider surfaces as a clean error instead of stalling the
+    # response stream indefinitely (incident 2026-07-29).
+    LLM_TIMEOUT: float = float(os.getenv("LLM_TIMEOUT", "90"))  # pylint: disable=invalid-name
+
+    # A call slower than this is logged at WARNING. Normal analysis calls
+    # completed in ~6s during the incident window, so this flags outliers
+    # well before they reach the timeout.
+    # pylint: disable-next=invalid-name
+    LLM_SLOW_CALL_THRESHOLD: float = float(os.getenv("LLM_SLOW_CALL_THRESHOLD", "20"))
+
+    # Retry budget for a single agent LLM call. Kept explicit because the
+    # provider SDK and litellm each have their own retry loop, and leaving
+    # both at their defaults multiplies the effective ceiling (measured: a
+    # 3s timeout took 10.8s to fail). Applied as the SDK-level retry count
+    # with litellm's outer loop disabled, so the worst case stays close to
+    # LLM_TIMEOUT rather than a multiple of it.
+    # pylint: disable-next=invalid-name
+    LLM_MAX_RETRIES: int = int(os.getenv("LLM_MAX_RETRIES", "1"))
+
     DB_MAX_DISTINCT: int = 100  # pylint: disable=invalid-name
     DB_UNIQUENESS_THRESHOLD: float = 0.5  # pylint: disable=invalid-name
     SHORT_MEMORY_LENGTH = 5  # Maximum number of questions to keep in short-term memory

api/core/text2sql.py

@@ -407,7 +407,12 @@ async def run_query(  # pylint: disable=too-many-locals,too-many-branches,too-ma
     agent_an = AnalysisAgent(
         queries_history, result_history, custom_api_key, custom_model,
     )
-    answer_an = agent_an.get_analysis(
+    # ``get_analysis`` is a synchronous LLM call. Running it directly here
+    # would block the event loop for its full duration, stalling every other
+    # in-flight request and preventing any stream from flushing bytes
+    # (incident 2026-07-29). Off-loop via a worker thread.
+    answer_an = await asyncio.to_thread(
+        agent_an.get_analysis,
         queries_history[-1], tables, db_description, instructions, memory_context,
         db_type, user_rules_spec,
     )
@@ -427,7 +432,8 @@ async def run_query(  # pylint: disable=too-many-locals,too-many-branches,too-ma
         follow_up_agent = FollowUpAgent(
             queries_history, result_history, custom_api_key, custom_model,
         )
-        follow_up = follow_up_agent.generate_follow_up_question(
+        follow_up = await asyncio.to_thread(
+            follow_up_agent.generate_follow_up_question,
             user_question=queries_history[-1],
             analysis_result=answer_an,
         )
@@ -536,7 +542,11 @@ def _run_sql(sql: str):
                     )
                 return loader_class.execute_sql_query(sql, db_url)
 
-            healing_result = healer.heal_and_execute(
+            # Same reasoning as ``get_analysis`` above, and worse here: this
+            # chains up to ``max_healing_attempts`` sequential LLM calls plus
+            # SQL execution, all synchronous.
+            healing_result = await asyncio.to_thread(
+                healer.heal_and_execute,
                 initial_sql=sql_query,
                 initial_error=str(exec_error),
                 execute_sql_func=_run_sql,
@@ -593,7 +603,8 @@ def _run_sql(sql: str):
             "message": f"Step {step_num}: Generating user-friendly response",
         }
 
-        user_readable_response = format_ai_response(
+        user_readable_response = await asyncio.to_thread(
+            format_ai_response,
             queries_history=queries_history,
             result_history=result_history,
             sql_query=sql_query,
@@ -755,7 +766,8 @@ async def run_confirmed(  # pylint: disable=too-many-locals,too-many-branches,to
         yield {"type": "reasoning_step",
                "message": f"Step {step_num}: Generating user-friendly response"}
 
-        user_readable_response = format_ai_response(
+        user_readable_response = await asyncio.to_thread(
+            format_ai_response,
             queries_history=queries_history or [question],
             result_history=None,
             sql_query=sql_query,

api/memory/graphiti_tool.py

@@ -27,7 +27,7 @@
 from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF
 
 
-from litellm import completion
+from api.agents.utils import run_completion
 
 
 def extract_embedding_model_name(full_model_name: str) -> str:
@@ -253,22 +253,22 @@ async def update_user_information(self, conversation: Dict[str, Any], history: T
                 """
         try:
 
-            if len(history[1]) == 0:
+            if not history[1]:
                 messages = [{"role": "user", "content": prompt}]
             else:
                 messages = []
                 for query, result in zip(history[0], history[1]):
                     messages.append({"role": "user", "content": query})
                     messages.append({"role": "assistant", "content": result})
             messages.append({"role": "user", "content": prompt})
-            response = completion(
-                model=Config.COMPLETION_MODEL,
-                messages=messages,
-                temperature=0.1
-            )
-
-            # Parse the direct text response (no JSON parsing needed)
-            content = response.choices[0].message.content.strip()
+            # Synchronous LLM call inside an async method that runs as a
+            # detached task: calling it directly would block the event loop
+            # and stall unrelated streaming responses. ``run_completion`` also
+            # applies the shared timeout and retry bounds.
+            content = (await asyncio.to_thread(
+                run_completion, messages, label="memory.user_summary",
+                temperature=0.1,
+            )).strip()
             query = """
             MATCH (u:Entity {name: $user_id})
             SET u.summary = $summary
@@ -277,6 +277,9 @@ async def update_user_information(self, conversation: Dict[str, Any], history: T
             await driver.execute_query(query, user_id=self.user_id, summary=content)
             return True
         except Exception as e:
+            # Previously swallowed silently, which hid a recurring failure on
+            # this path entirely (incident 2026-07-29).
+            logging.error("Error updating user information: %s", e)
             return False
 
     async def add_new_memory(self, conversation: Dict[str, Any], history: Tuple[List[str], List[str]]) -> bool:
@@ -733,22 +736,20 @@ async def summarize_conversation(self, conversation: Dict[str, Any], history: Li
 
         try:
 
-            if len(history[1]) == 0:
+            if not history[1]:
                 messages = [{"role": "user", "content": prompt}]
             else:
                 messages = []
                 for query, result in zip(history[0], history[1]):
                     messages.append({"role": "user", "content": query})
                     messages.append({"role": "assistant", "content": result})
             messages.append({"role": "user", "content": prompt})
-            response = completion(
-                model=Config.COMPLETION_MODEL,
-                messages=messages,
-                temperature=0.1
-            )
-
-            # Parse the direct text response (no JSON parsing needed)
-            content = response.choices[0].message.content.strip()
+            # Same reasoning as ``update_user_information`` above: off-loop,
+            # with the shared timeout and retry bounds.
+            content = (await asyncio.to_thread(
+                run_completion, messages, label="memory.conversation_summary",
+                temperature=0.1,
+            )).strip()
             return {
                 "database_summary": content
             }
@@ -769,7 +770,10 @@ def __init__(self):
 
         self.api_key = os.getenv('AZURE_API_KEY')
         self.endpoint = os.getenv('AZURE_API_BASE') 
-        self.api_version = os.getenv('AZURE_API_VERSION', '2024-02-01')
+        # Graphiti's OpenAI client uses the Responses API, which requires
+        # api-version 2025-03-01-preview or later. Older values make every
+        # episode write fail with HTTP 400 (incident 2026-07-29).
+        self.api_version = os.getenv('AZURE_API_VERSION', '2025-03-01-preview')
         self.model_choice = "gpt-4.1"  # Use the model name directly
 
         # Extract just the model name without provider prefix for Graphiti

api/routes/database.py

@@ -5,6 +5,7 @@
 
 from api.auth.user_management import token_required
 from api.core.schema_loader import load_database
+from api.routes.streaming import STREAM_HEADERS, with_keepalive
 from api.routes.tokens import UNAUTHORIZED_RESPONSE
 
 database_router = APIRouter(tags=["Database Connection"])
@@ -30,4 +31,8 @@ async def connect_database(request: Request, db_request: DatabaseConnectionReque
     Requires authentication.
     """
     generator = await load_database(db_request.url, request.state.user_id)
-    return StreamingResponse(generator, media_type="application/json")
+    return StreamingResponse(
+        with_keepalive(generator),
+        media_type="application/json",
+        headers=STREAM_HEADERS,
+    )