55import json
66import logging
77import time
8+ from functools import partial
89from typing import (
910 Annotated ,
1011 Any ,
1112 Optional ,
1213 Union ,
1314)
1415
16+ import anyio
1517from fastapi import (
1618 Body ,
1719 Path ,
5961# Keep OpenAI as a fallback option
6062try :
6163 import openai
64+ from openai import AsyncOpenAI
6265except ImportError :
6366 openai = None # type: ignore[assignment]
67+ AsyncOpenAI = None # type: ignore[assignment,misc]
6468
6569log = logging .getLogger (__name__ )
6670
@@ -152,9 +156,9 @@ async def query(
152156 job = None
153157 if job_id :
154158 # Job-based chat - check for existing responses (unless regenerate requested)
155- job = self .job_manager .get_accessible_job ( trans , job_id )
159+ job = await anyio . to_thread . run_sync ( partial ( self .job_manager .get_accessible_job , trans , job_id ) )
156160 if job and not regenerate :
157- existing_response = self .chat_manager .get ( trans , job .id )
161+ existing_response = await anyio . to_thread . run_sync ( partial ( self .chat_manager .get , trans , job .id ) )
158162 if existing_response and existing_response .messages [0 ]:
159163 return ChatResponse (
160164 response = existing_response .messages [0 ].message ,
@@ -176,7 +180,9 @@ async def query(
176180
177181 # If we have an exchange_id, ALWAYS load conversation history from database (source of truth)
178182 if exchange_id :
179- db_history = self .chat_manager .get_chat_history (trans , exchange_id , format_for_pydantic_ai = False )
183+ db_history = await anyio .to_thread .run_sync (
184+ partial (self .chat_manager .get_chat_history , trans , exchange_id , format_for_pydantic_ai = False )
185+ )
180186 if db_history :
181187 full_context ["conversation_history" ] = db_history
182188 else :
@@ -197,13 +203,15 @@ async def query(
197203 self ._ensure_ai_configured ()
198204 # For legacy, use context_type from query_context if it exists
199205 context_type = query_context .get ("context_type" ) if isinstance (query_context , dict ) else None
200- answer = self ._get_ai_response (query_text , trans , context_type )
206+ answer = await self ._get_ai_response (query_text , trans , context_type )
201207 result ["response" ] = answer
202208
203209 # Save chat exchange to database
204210 if job :
205211 # Job-based chat
206- exchange = self .chat_manager .create (trans , job .id , str (result ["response" ]))
212+ exchange = await anyio .to_thread .run_sync (
213+ partial (self .chat_manager .create , trans , job .id , str (result ["response" ]))
214+ )
207215 result ["exchange_id" ] = exchange .id
208216 elif trans .user :
209217 # Use the exchange_id we already extracted at the beginning
@@ -217,7 +225,9 @@ async def query(
217225 "agent_response" : agent_resp .model_dump () if agent_resp else None ,
218226 }
219227 message_content = json .dumps (conversation_data )
220- self .chat_manager .add_message (trans , exchange_id , message_content )
228+ await anyio .to_thread .run_sync (
229+ partial (self .chat_manager .add_message , trans , exchange_id , message_content )
230+ )
221231 result ["exchange_id" ] = exchange_id
222232 else :
223233 # Create new exchange for first message
@@ -227,7 +237,9 @@ async def query(
227237 "response" : result .get ("response" , "" ),
228238 "agent_response" : agent_resp .model_dump () if agent_resp else None ,
229239 }
230- exchange = self .chat_manager .create_general_chat (trans , query_text , storable_result , agent_type )
240+ exchange = await anyio .to_thread .run_sync (
241+ partial (self .chat_manager .create_general_chat , trans , query_text , storable_result , agent_type )
242+ )
231243 result ["exchange_id" ] = exchange .id
232244
233245 result ["processing_time" ] = time .time () - start_time
@@ -421,7 +433,7 @@ def _ensure_ai_configured(self):
421433 if self .config .ai_api_key is None :
422434 raise ConfigurationError ("AI API key is not configured for this instance." )
423435
424- def _get_ai_response (self , query : str , trans : ProvidesUserContext , context_type : Optional [str ] = None ) -> str :
436+ async def _get_ai_response (self , query : str , trans : ProvidesUserContext , context_type : Optional [str ] = None ) -> str :
425437 """Get response from AI using pydantic-ai Agent"""
426438 system_prompt = self ._get_system_prompt ()
427439 username = trans .user .username if trans .user else "Anonymous User"
@@ -432,8 +444,8 @@ def _get_ai_response(self, query: str, trans: ProvidesUserContext, context_type:
432444 full_system_prompt = f"{ system_prompt } \n \n You will address the user as { username } "
433445 agent : Agent [None , str ] = Agent (model_name , system_prompt = full_system_prompt )
434446
435- # Get response from the agent
436- result = agent .run_sync (query )
447+ # Get response from the agent (async)
448+ result = await agent .run (query )
437449 return result .output
438450 except UnexpectedModelBehavior as e :
439451 log .error (f"Unexpected model behavior: { e } " )
@@ -442,10 +454,10 @@ def _get_ai_response(self, query: str, trans: ProvidesUserContext, context_type:
442454 log .error (f"Error using pydantic-ai Agent: { e } " )
443455 # Try fallback to direct OpenAI if available
444456 if openai is not None :
445- return self ._call_openai_directly (query , system_prompt , username )
457+ return await self ._call_openai_directly (query , system_prompt , username )
446458 raise
447459
448- def _call_openai_directly (self , query : str , system_prompt : str , username : str ) -> str :
460+ async def _call_openai_directly (self , query : str , system_prompt : str , username : str ) -> str :
449461 """Direct OpenAI API call as fallback"""
450462 try :
451463 messages : list [dict [str , str ]] = [
@@ -456,7 +468,8 @@ def _call_openai_directly(self, query: str, system_prompt: str, username: str) -
456468 },
457469 {"role" : "user" , "content" : query },
458470 ]
459- response = openai .chat .completions .create (
471+ client = AsyncOpenAI ()
472+ response = await client .chat .completions .create (
460473 model = self .config .ai_model ,
461474 messages = messages , # type: ignore[arg-type]
462475 )
0 commit comments