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

.env.example

@@ -75,6 +75,34 @@ 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 budget for one LLM call end to end, in seconds (default 90).
+# This is the total, not per attempt: the per-attempt timeout handed to the
+# provider is this divided by LLM_MAX_RETRIES + 1, so retries cannot push the
+# real ceiling past it. A hung provider surfaces as an error within the budget
+# 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 +128,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()