Open AI: Running Threads synchronously (#143)

* first stab at adding sync api

* refactoring few things

* added testcases

* fixing few testcases

* added few more testcases

Author: AkhileshNegi

backend/app/api/routes/threads.py

@@ -29,13 +29,89 @@ def send_callback(callback_url: str, data: dict):
         return False
 
 
+def handle_openai_error(e: openai.OpenAIError) -> str:
+    """Extract error message from OpenAI error."""
+    if isinstance(e.body, dict) and "message" in e.body:
+        return e.body["message"]
+    return str(e)
+
+
+def validate_thread(client: OpenAI, thread_id: str) -> tuple[bool, str]:
+    """Validate if a thread exists and has no active runs."""
+    if not thread_id:
+        return True, None
+
+    try:
+        runs = client.beta.threads.runs.list(thread_id=thread_id)
+        if runs.data and len(runs.data) > 0:
+            latest_run = runs.data[0]
+            if latest_run.status in ["queued", "in_progress", "requires_action"]:
+                return (
+                    False,
+                    f"There is an active run on this thread (status: {latest_run.status}). Please wait for it to complete.",
+                )
+        return True, None
+    except openai.OpenAIError:
+        return False, f"Invalid thread ID provided {thread_id}"
+
+
+def setup_thread(client: OpenAI, request: dict) -> tuple[bool, str]:
+    """Set up thread and add message, either creating new or using existing."""
+    thread_id = request.get("thread_id")
+
+    if thread_id:
+        try:
+            client.beta.threads.messages.create(
+                thread_id=thread_id, role="user", content=request["question"]
+            )
+            return True, None
+        except openai.OpenAIError as e:
+            return False, handle_openai_error(e)
+    else:
+        try:
+            thread = client.beta.threads.create()
+            client.beta.threads.messages.create(
+                thread_id=thread.id, role="user", content=request["question"]
+            )
+            request["thread_id"] = thread.id
+            return True, None
+        except openai.OpenAIError as e:
+            return False, handle_openai_error(e)
+
+
+def process_message_content(message_content: str, remove_citation: bool) -> str:
+    """Process message content, optionally removing citations."""
+    if remove_citation:
+        return re.sub(r"【\d+(?::\d+)?†[^】]*】", "", message_content)
+    return message_content
+
+
+def get_additional_data(request: dict) -> dict:
+    """Extract additional data from request, excluding specific keys."""
+    return {
+        k: v
+        for k, v in request.items()
+        if k not in {"question", "assistant_id", "callback_url", "thread_id"}
+    }
+
+
+def create_success_response(request: dict, message: str) -> APIResponse:
+    """Create a success response with the given message and request data."""
+    additional_data = get_additional_data(request)
+    return APIResponse.success_response(
+        data={
+            "status": "success",
+            "message": message,
+            "thread_id": request["thread_id"],
+            "endpoint": getattr(request, "endpoint", "some-default-endpoint"),
+            **additional_data,
+        }
+    )
+
+
 def process_run(request: dict, client: OpenAI):
-    """
-    Background task to run create_and_poll, then send the callback with the result.
-    This function is run in the background after we have already returned an initial response.
-    """
+    """Process a run and send callback with results."""
     try:
-        # Start the run
         run = client.beta.threads.runs.create_and_poll(
             thread_id=request["thread_id"],
             assistant_id=request["assistant_id"],
@@ -45,46 +121,19 @@ def process_run(request: dict, client: OpenAI):
             messages = client.beta.threads.messages.list(thread_id=request["thread_id"])
             latest_message = messages.data[0]
             message_content = latest_message.content[0].text.value
-
-            remove_citation = request.get("remove_citation", False)
-
-            if remove_citation:
-                message = re.sub(r"【\d+(?::\d+)?†[^】]*】", "", message_content)
-            else:
-                message = message_content
-
-            # Update the data dictionary with additional fields from the request, excluding specific keys
-            additional_data = {
-                k: v
-                for k, v in request.items()
-                if k not in {"question", "assistant_id", "callback_url", "thread_id"}
-            }
-            callback_response = APIResponse.success_response(
-                data={
-                    "status": "success",
-                    "message": message,
-                    "thread_id": request["thread_id"],
-                    "endpoint": getattr(request, "endpoint", "some-default-endpoint"),
-                    **additional_data,
-                }
+            message = process_message_content(
+                message_content, request.get("remove_citation", False)
             )
+            callback_response = create_success_response(request, message)
         else:
             callback_response = APIResponse.failure_response(
                 error=f"Run failed with status: {run.status}"
             )
 
-        # Send callback with results
         send_callback(request["callback_url"], callback_response.model_dump())
 
     except openai.OpenAIError as e:
-        # Handle any other OpenAI API errors
-        if isinstance(e.body, dict) and "message" in e.body:
-            error_message = e.body["message"]
-        else:
-            error_message = str(e)
-
-        callback_response = APIResponse.failure_response(error=error_message)
-
+        callback_response = APIResponse.failure_response(error=handle_openai_error(e))
         send_callback(request["callback_url"], callback_response.model_dump())
 
 
@@ -95,54 +144,20 @@ async def threads(
     _session: Session = Depends(get_db),
     _current_user: UserOrganization = Depends(get_current_user_org),
 ):
-    """
-    Accepts a question, assistant_id, callback_url, and optional thread_id from the request body.
-    Returns an immediate "processing" response, then continues to run create_and_poll in background.
-    Once completed, calls send_callback with the final result.
-    """
+    """Asynchronous endpoint that processes requests in background."""
     client = OpenAI(api_key=settings.OPENAI_API_KEY)
 
-    # Use get method to safely access thread_id
-    thread_id = request.get("thread_id")
+    # Validate thread
+    is_valid, error_message = validate_thread(client, request.get("thread_id"))
+    if not is_valid:
+        return APIResponse.failure_response(error=error_message)
 
-    # 1. Validate or check if there's an existing thread with an in-progress run
-    if thread_id:
-        try:
-            runs = client.beta.threads.runs.list(thread_id=thread_id)
-            # Get the most recent run (first in the list) if any
-            if runs.data and len(runs.data) > 0:
-                latest_run = runs.data[0]
-                if latest_run.status in ["queued", "in_progress", "requires_action"]:
-                    return APIResponse.failure_response(
-                        error=f"There is an active run on this thread (status: {latest_run.status}). Please wait for it to complete."
-                    )
-        except openai.OpenAIError:
-            # Handle invalid thread ID
-            return APIResponse.failure_response(
-                error=f"Invalid thread ID provided {thread_id}"
-            )
+    # Setup thread
+    is_success, error_message = setup_thread(client, request)
+    if not is_success:
+        return APIResponse.failure_response(error=error_message)
 
-        # Use existing thread
-        client.beta.threads.messages.create(
-            thread_id=thread_id, role="user", content=request["question"]
-        )
-    else:
-        try:
-            # Create new thread
-            thread = client.beta.threads.create()
-            client.beta.threads.messages.create(
-                thread_id=thread.id, role="user", content=request["question"]
-            )
-            request["thread_id"] = thread.id
-        except openai.OpenAIError as e:
-            # Handle any other OpenAI API errors
-            if isinstance(e.body, dict) and "message" in e.body:
-                error_message = e.body["message"]
-            else:
-                error_message = str(e)
-            return APIResponse.failure_response(error=error_message)
-
-    # 2. Send immediate response to complete the API call
+    # Send immediate response
     initial_response = APIResponse.success_response(
         data={
             "status": "processing",
@@ -152,8 +167,50 @@ async def threads(
         }
     )
 
-    # 3. Schedule the background task to run create_and_poll and send callback
+    # Schedule background task
     background_tasks.add_task(process_run, request, client)
 
-    # 4. Return immediately so the client knows we've accepted the request
     return initial_response
+
+
+@router.post("/threads/sync")
+async def threads_sync(
+    request: dict,
+    _session: Session = Depends(get_db),
+    _current_user: UserOrganization = Depends(get_current_user_org),
+):
+    """Synchronous endpoint that processes requests immediately."""
+    client = OpenAI(api_key=settings.OPENAI_API_KEY)
+
+    # Validate thread
+    is_valid, error_message = validate_thread(client, request.get("thread_id"))
+    if not is_valid:
+        return APIResponse.failure_response(error=error_message)
+
+    # Setup thread
+    is_success, error_message = setup_thread(client, request)
+    if not is_success:
+        return APIResponse.failure_response(error=error_message)
+
+    try:
+        # Process run
+        run = client.beta.threads.runs.create_and_poll(
+            thread_id=request["thread_id"],
+            assistant_id=request["assistant_id"],
+        )
+
+        if run.status == "completed":
+            messages = client.beta.threads.messages.list(thread_id=request["thread_id"])
+            latest_message = messages.data[0]
+            message_content = latest_message.content[0].text.value
+            message = process_message_content(
+                message_content, request.get("remove_citation", False)
+            )
+            return create_success_response(request, message)
+        else:
+            return APIResponse.failure_response(
+                error=f"Run failed with status: {run.status}"
+            )
+
+    except openai.OpenAIError as e:
+        return APIResponse.failure_response(error=handle_openai_error(e))