Remove crew, replace with custom tool-calling (#243)

* wip on moving away from crew

* wip

lastFile:llm-service/app/services/query/agents/crewai_querier.py

* drop databases

lastFile:llm-service/app/services/query/agents/crewai_querier.py

* WIP basic tool calling and minor refactoring of retriever tool

lastFile:llm-service/app/services/query/tools/retriever.py

* not working usage of openai agent

* basic streaming example

* wip on openai tool using

* provide summaries to the retrieval tool

* added source nodes to OpenAI agent stream

* minor change to prompt

* wip on supporting bedrock

* WIP streaming for non open ai models

* add conditional for openai

* support for non openai agents and refactor

* minor fix to return types

* remove crew

* minor change to check for empty response

* fix mypy, ruff issues

* Revert "bump mui packages that seemed to address issues with x-charts"

This reverts commit 864b770.

* Revert "update hide legend"

This reverts commit b4db0dd.

* remove the last crew bits

* added verbosity for non openai agents

* added missing print dashes

* better error handling for bedrock model availability

* fix mypy

* fix writing source nodes

---------

Co-authored-by: Elijah Williams <ewilliams@cloudera.com>
Co-authored-by: Michael Liu <mliu@cloudera.com>
Co-authored-by: Baasit Sharief <baasitsharief@gmail.com>

Author: 4 people

llm-service/app/routers/index/sessions/__init__.py

@@ -63,8 +63,8 @@
 from ....services.chat_history.paginator import paginate
 from ....services.metadata_apis import session_metadata_api
 from ....services.mlflow import rating_mlflow_log_metric, feedback_mlflow_log_table
-from ....services.query.agents.crewai_querier import poison_pill
-from ....services.query.crew_events import CrewEvent
+from ....services.query.agents.tool_calling_querier import poison_pill
+from ....services.query.chat_events import ToolEvent
 from ....services.session import rename_session
 
 logger = logging.getLogger(__name__)
@@ -258,15 +258,15 @@ def stream_chat_completion(
     session = session_metadata_api.get_session(session_id, user_name=origin_remote_user)
     configuration = request.configuration or RagPredictConfiguration()
 
-    crew_events_queue: queue.Queue[CrewEvent] = queue.Queue()
+    tool_events_queue: queue.Queue[ToolEvent] = queue.Queue()
     # Create a cancellation event to signal when the client disconnects
     cancel_event = threading.Event()
 
-    def crew_callback(chat_future: Future[Any]) -> Generator[str, None, None]:
+    def tools_callback(chat_future: Future[Any]) -> Generator[str, None, None]:
         while True:
             # Check if client has disconnected
             if cancel_event.is_set():
-                logger.info("Client disconnected, stopping crew callback")
+                logger.info("Client disconnected, stopping tool callback")
                 # Try to cancel the future if it's still running
                 if not chat_future.done():
                     chat_future.cancel()
@@ -276,14 +276,14 @@ def crew_callback(chat_future: Future[Any]) -> Generator[str, None, None]:
                 raise e
 
             try:
-                event_data = crew_events_queue.get(block=True, timeout=1.0)
+                event_data = tool_events_queue.get(block=True, timeout=1.0)
                 if event_data.type == poison_pill:
                     break
                 event_json = json.dumps({"event": event_data.model_dump()})
                 yield f"data: {event_json}\n\n"
             except queue.Empty:
                 # Send a heartbeat event every second to keep the connection alive
-                heartbeat = CrewEvent(
+                heartbeat = ToolEvent(
                     type="event", name="Processing", timestamp=time.time()
                 )
                 event_json = json.dumps({"event": heartbeat.model_dump()})
@@ -303,27 +303,28 @@ def generate_stream() -> Generator[str, None, None]:
                 query=request.query,
                 configuration=configuration,
                 user_name=origin_remote_user,
-                crew_events_queue=crew_events_queue,
+                tool_events_queue=tool_events_queue,
             )
 
-            # Yield from crew_callback, which will check for cancellation
-            yield from crew_callback(future)
+            # Yield from tools_callback, which will check for cancellation
+            yield from tools_callback(future)
 
             # If we get here and the cancel_event is set, the client has disconnected
             if cancel_event.is_set():
                 logger.info("Client disconnected, not processing results")
                 return
 
             first_message = True
-            for response in future.result():
+            stream = future.result()
+            for response in stream:
                 # Check for cancellation between each response
                 if cancel_event.is_set():
                     logger.info("Client disconnected during result processing")
                     break
 
                 # send an initial message to let the client know the response stream is starting
                 if first_message:
-                    done = CrewEvent(type="done", name="done", timestamp=time.time())
+                    done = ToolEvent(type="done", name="done", timestamp=time.time())
                     event_json = json.dumps({"event": done.model_dump()})
                     yield f"data: {event_json}\n\n"
                     first_message = False

llm-service/app/services/caii/caii.py

@@ -131,7 +131,6 @@ def get_llm(
         base_url=api_base,
         model=model,
         http_client=http_client
-        # api_base=api_base, # todo: figure out how to integrate with Crew models
     )
 
 

llm-service/app/services/chat/streaming_chat.py

@@ -59,12 +59,12 @@
 from app.services.metadata_apis.session_metadata_api import Session
 from app.services.mlflow import record_direct_llm_mlflow_run
 from app.services.query import querier
-from app.services.query.agents.crewai_querier import poison_pill
+from app.services.query.agents.tool_calling_querier import poison_pill
 from app.services.query.chat_engine import (
     FlexibleContextChatEngine,
     build_flexible_chat_engine,
 )
-from app.services.query.crew_events import CrewEvent
+from app.services.query.chat_events import ToolEvent
 from app.services.query.querier import (
     build_retriever,
 )
@@ -76,7 +76,7 @@ def stream_chat(
     query: str,
     configuration: RagPredictConfiguration,
     user_name: Optional[str],
-    crew_events_queue: Queue[CrewEvent],
+    tool_events_queue: Queue[ToolEvent],
 ) -> Generator[ChatResponse, None, None]:
     query_configuration = QueryConfiguration(
         top_k=session.response_chunks,
@@ -99,14 +99,12 @@ def stream_chat(
     if not query_configuration.use_tool_calling and (
         len(session.data_source_ids) == 0 or total_data_sources_size == 0
     ):
-        # put a poison pill in the queue to stop the crew events stream
-        crew_events_queue.put(CrewEvent(type=poison_pill, name="no-op"))
-        return _stream_direct_llm_chat(
-            session, response_id, query, user_name, crew_events_queue
-        )
+        # put a poison pill in the queue to stop the tool events stream
+        tool_events_queue.put(ToolEvent(type=poison_pill, name="no-op"))
+        return _stream_direct_llm_chat(session, response_id, query, user_name)
 
     condensed_question, streaming_chat_response = build_streamer(
-        crew_events_queue, query, query_configuration, session
+        tool_events_queue, query, query_configuration, session
     )
     return _run_streaming_chat(
         session,
@@ -125,10 +123,11 @@ def _run_streaming_chat(
     query: str,
     query_configuration: QueryConfiguration,
     user_name: Optional[str],
+    streaming_chat_response: StreamingAgentChatResponse,
     condensed_question: Optional[str] = None,
-    streaming_chat_response: StreamingAgentChatResponse = None,
 ) -> Generator[ChatResponse, None, None]:
     response: ChatResponse = ChatResponse(message=ChatMessage(content=query))
+
     if streaming_chat_response.chat_stream:
         for response in streaming_chat_response.chat_stream:
             response.additional_kwargs["response_id"] = response_id
@@ -152,7 +151,7 @@ def _run_streaming_chat(
 
 
 def build_streamer(
-    crew_events_queue: Queue[CrewEvent],
+    chat_events_queue: Queue[ToolEvent],
     query: str,
     query_configuration: QueryConfiguration,
     session: Session,
@@ -181,9 +180,8 @@ def build_streamer(
         query,
         query_configuration,
         chat_messages,
-        crew_events_queue=crew_events_queue,
+        tool_events_queue=chat_events_queue,
         session=session,
-        retriever=retriever,
     )
     return condensed_question, streaming_chat_response
 
@@ -193,7 +191,6 @@ def _stream_direct_llm_chat(
     response_id: str,
     query: str,
     user_name: Optional[str],
-    queue: Queue[CrewEvent],
 ) -> Generator[ChatResponse, None, None]:
     record_direct_llm_mlflow_run(response_id, session, user_name)
 

llm-service/app/services/models/providers/bedrock.py

@@ -44,6 +44,7 @@
 import requests
 from botocore.auth import SigV4Auth
 from botocore.awsrequest import AWSRequest
+from fastapi import HTTPException
 from llama_index.embeddings.bedrock import BedrockEmbedding
 from llama_index.llms.bedrock_converse import BedrockConverse
 from llama_index.llms.bedrock_converse.utils import BEDROCK_MODELS
@@ -153,22 +154,33 @@ def get_aws_responses(
             raise_for_http_error(response)
             return cast(dict[str, Any], response.json())
 
-        responses: list[dict[str, Any] | None] = [None for _ in aws_requests]
+        responses: list[dict[str, Any] | None] = []
         with concurrent.futures.ThreadPoolExecutor() as executor:
-            future_to_index = {
-                executor.submit(get_aws_responses, url, headers): idx
-                for idx, (url, headers) in enumerate(aws_requests)
-            }
-            for future in concurrent.futures.as_completed(future_to_index):
-                idx = future_to_index[future]
+            results = executor.map(
+                lambda url_and_headers: get_aws_responses(*url_and_headers),
+                aws_requests,
+            )
+            while True:
                 try:
-                    responses[idx] = future.result()
-                except Exception:
+                    result = next(results)
+                    responses.append(result)
+                except StopIteration:
+                    break
+                except HTTPException as e:
+                    model_id = str(e).split("/")[-1]
                     logger.exception(
-                        "Error fetching data for model %s", models[idx]["modelId"]
+                        "Error fetching data for model %s",
+                        model_id,
                     )
-                    responses[idx] = None
-
+                    responses.append(None)
+                    continue
+                except Exception as e:
+                    logger.exception(
+                        "Unexpected error fetching data: %s",
+                        e,
+                    )
+                    responses.append(None)
+                    continue
         for model, model_data in zip(models, responses):
             if model_data:
                 if model_data["entitlementAvailability"] == "AVAILABLE":