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

.env.example

@@ -75,6 +75,32 @@ 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
+#
+# Bounds for executing a user query against the target database. Offloading
+# execution to a thread keeps the event loop free, but only these bound how
+# long the query itself may run (a thread blocked in a socket read cannot be
+# cancelled from Python). Seconds.
+# DB_CONNECT_TIMEOUT=10
+# DB_STATEMENT_TIMEOUT=60
+#
+# Schema introspection gets a larger deadline (it is metadata work over a whole
+# database) and a cap on how many may occupy worker threads at once, since that
+# executor is shared with every other offloaded call.
+# DB_SCHEMA_TIMEOUT=300
+# DB_SCHEMA_CONCURRENCY=2
+
 # OpenAI - uses openai/gpt-4.1 and openai/text-embedding-ada-002
 # OPENAI_API_KEY=your_openai_api_key
 
@@ -100,7 +126,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

@@ -4,7 +4,9 @@
 """
 
 import os
+import time
 import logging
+import math
 import dataclasses
 from typing import Union
 
@@ -40,30 +42,64 @@ def __init__(self, model_name: str, config: dict = None):
         self.model_name = model_name
         self.config = config
 
+    def _embedding_kwargs(self) -> dict:
+        """Timeout and retry bounds, matching the completion path.
+
+        These are blocking network calls, so an unbounded one pins whichever
+        thread runs it. ``timeout`` is per attempt, so the retry budget is
+        pinned too or the effective ceiling becomes a multiple of it.
+        """
+        return {
+            "timeout": Config.LLM_TIMEOUT,
+            "max_retries": Config.LLM_MAX_RETRIES,
+            "num_retries": 0,
+        }
+
     def embed(self, text: Union[str, list]) -> list:
         """
         Get the embeddings of the text
 
+        Blocking: call via ``api.embeddings.embed_off_loop`` from async code.
+
         Args:
             text (str|list): The text(s) to embed
 
         Returns:
             list: The embeddings of the text
 
         """
-        embeddings = embedding(model=self.model_name, input=text)
+        started = time.monotonic()
+        try:
+            embeddings = embedding(
+                model=self.model_name, input=text, **self._embedding_kwargs()
+            )
+        except Exception:
+            logging.warning(
+                "embed_call model=%s duration=%.2fs outcome=error",
+                self.model_name, time.monotonic() - started,
+            )
+            raise
+        logging.info(
+            "embed_call model=%s duration=%.2fs outcome=ok",
+            self.model_name, time.monotonic() - started,
+        )
         embeddings = [embedding["embedding"] for embedding in embeddings.data]
         return embeddings
 
     def get_vector_size(self) -> int:
         """
         Get the size of the vector
 
+        Blocking: call via ``api.embeddings.vector_size_off_loop`` from async
+        code.
+
         Returns:
             int: The size of the vector
 
         """
-        response = embedding(input=["Hello World"], model=self.model_name)
+        response = embedding(
+            input=["Hello World"], model=self.model_name, **self._embedding_kwargs()
+        )
         size = len(response.data[0]["embedding"])
         return size
 
@@ -77,8 +113,38 @@ def _with_prefix(model: str, provider: str) -> str:
 SUPPORTED_VENDORS = ("openai", "anthropic", "gemini", "azure", "ollama", "cohere")
 
 
+def _positive_env(name: str, default: str, cast=int):
+    """Read a timeout-style env var, rejecting values that disable the bound.
+
+    Zero is not a harmless "unset" here: PostgreSQL treats a 0 timeout as
+    "no limit", which removes the safeguard entirely, and PyMySQL raises at
+    query time on a 0 socket timeout. Fail at startup with a clear message
+    rather than silently losing the protection.
+    """
+    raw = os.getenv(name, default)
+    try:
+        value = cast(raw)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            f"{name} must be a positive number (got {raw!r})"
+        ) from exc
+    # ``nan`` and ``inf`` are floats that pass a ``<= 0`` test: nan compares
+    # False against everything, and inf is a deadline that never expires.
+    if not math.isfinite(value):
+        raise ValueError(
+            f"{name} must be a finite number (got {raw!r}); nan and inf are not "
+            "usable deadlines"
+        )
+    if value <= 0:
+        raise ValueError(
+            f"{name} must be greater than 0 (got {raw!r}); a zero or negative "
+            "timeout disables the safeguard it exists to provide"
+        )
+    return value
+
+
 @dataclasses.dataclass
-class Config:
+class Config:  # pylint: disable=too-many-instance-attributes
     """
     Configuration class for the text2sql module.
     """
@@ -133,6 +199,54 @@ 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 = _positive_env("LLM_TIMEOUT", "90", float)  # 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 = _positive_env(
+        "LLM_SLOW_CALL_THRESHOLD", "20", float
+    )
+
+    # 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.
+    # Zero is valid here (it means "no retry", a strict ceiling); negative is
+    # not.
+    # pylint: disable-next=invalid-name
+    LLM_MAX_RETRIES: int = max(0, int(os.getenv("LLM_MAX_RETRIES", "1")))
+
+    # Bounds for user-query execution against the target database. Offloading
+    # execution to a thread stops a slow query from blocking other requests,
+    # but nothing bounds how long the query itself runs without these.
+    # pylint: disable-next=invalid-name
+    DB_CONNECT_TIMEOUT: int = _positive_env("DB_CONNECT_TIMEOUT", "10")
+    # pylint: disable-next=invalid-name
+    DB_STATEMENT_TIMEOUT: int = _positive_env("DB_STATEMENT_TIMEOUT", "60")
+
+    # Schema introspection gets its own, larger deadline: it is metadata work
+    # over a whole database, so the user-query ceiling is too tight, but it
+    # still needs a bound. Cancelling the awaiting task does not stop the
+    # driver call, so without this a stalled database holds both a session and
+    # a worker thread until it decides to answer.
+    # pylint: disable-next=invalid-name
+    DB_SCHEMA_TIMEOUT: int = _positive_env("DB_SCHEMA_TIMEOUT", "300")
+
+    # How many schema introspections may occupy worker threads at once. The
+    # default executor is shared with every other offloaded call (LLM,
+    # embedding, user SQL), so unbounded schema work on a stalled database
+    # could starve all of it.
+    # pylint: disable-next=invalid-name
+    DB_SCHEMA_CONCURRENCY: int = _positive_env("DB_SCHEMA_CONCURRENCY", "2")
+
     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