From 28dcbfd5b34bd7ab07dd1a95f61ce21a0a3dba9f Mon Sep 17 00:00:00 2001 From: Amitesh Gupta Date: Fri, 10 Jul 2026 13:02:45 +0530 Subject: [PATCH] fix: Fixed all 4 bugs from issue #243: initialized tool_calls before Addresses #243. --- backend/src/agents/retriever_rag.py | 7 +- backend/src/api/main.py | 1 + backend/src/api/routers/conversations.py | 21 +++-- .../tests/test_api_conversations_streaming.py | 61 +++++++++++++ backend/tests/test_process_html.py | 8 +- .../tests/test_retriever_graph_simplified.py | 86 +++++++++++++++++++ 6 files changed, 171 insertions(+), 13 deletions(-) diff --git a/backend/src/agents/retriever_rag.py b/backend/src/agents/retriever_rag.py index a90260cf..df21920c 100644 --- a/backend/src/agents/retriever_rag.py +++ b/backend/src/agents/retriever_rag.py @@ -107,7 +107,9 @@ def rag_initialize(self) -> None: self.tool_descriptions = "" for tool in self.tools: text_desc = render_text_description([tool]) - text_desc.replace("(query: str) -> Tuple[str, list[str], list[str]]", " ") + text_desc = text_desc.replace( + "(query: str) -> Tuple[str, list[str], list[str]]", " " + ) self.tool_descriptions += text_desc + "\n\n" def rag_agent(self, state: AgentState) -> dict[str, list[Any]]: @@ -159,12 +161,13 @@ def rag_agent(self, state: AgentState) -> dict[str, list[Any]]: ) return {"tools": []} + tool_calls: list = [] if "tool_names" in str(response): tool_calls = response.get("tool_names", []) # type: ignore for tool in tool_calls: if tool not in self.tool_names: logging.warning(f"Tool {tool} not found in tool list.") - tool_calls.remove(tool) + tool_calls = [tool for tool in tool_calls if tool in self.tool_names] else: logging.warning(str(response)) logging.warning("Tool selection failed. Returning empty tool list.") diff --git a/backend/src/api/main.py b/backend/src/api/main.py index 3ed04a34..ec7e4d5a 100644 --- a/backend/src/api/main.py +++ b/backend/src/api/main.py @@ -16,6 +16,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: init_database() logger.info("Starting graph initialization in background...") from .routers.conversations import start_graph_init + start_graph_init() yield logger.info("Shutting down...") diff --git a/backend/src/api/routers/conversations.py b/backend/src/api/routers/conversations.py index f6db750c..f005f33b 100644 --- a/backend/src/api/routers/conversations.py +++ b/backend/src/api/routers/conversations.py @@ -1,6 +1,7 @@ import os import logging import threading +from collections import OrderedDict from dotenv import load_dotenv from typing import Any, Optional @@ -257,7 +258,15 @@ async def ready() -> dict[str, str]: return {"status": "initializing"} -chat_history: dict[UUID, list[dict[str, str]]] = {} +MAX_IN_MEMORY_CONVERSATIONS = int(os.getenv("MAX_IN_MEMORY_CONVERSATIONS", "1000")) +chat_history: OrderedDict[UUID, list[dict[str, str]]] = OrderedDict() + + +def add_conversation_to_history(conversation_uuid: UUID) -> None: + """Registers a new conversation, evicting the oldest one if over capacity.""" + chat_history[conversation_uuid] = [] + if len(chat_history) > MAX_IN_MEMORY_CONVERSATIONS: + chat_history.popitem(last=False) def get_history_str(db: Session | None, conversation_uuid: UUID | None) -> str: @@ -314,7 +323,7 @@ async def get_agent_response( conversation_uuid = uuid4() if conversation_uuid not in chat_history: - chat_history[conversation_uuid] = [] + add_conversation_to_history(conversation_uuid) inputs = { "messages": [ @@ -327,7 +336,9 @@ async def get_agent_response( if graph is not None and graph.graph is not None: output = list(graph.graph.stream(inputs, stream_mode="updates")) else: - raise HTTPException(status_code=503, detail="Graph is still initializing. Please retry shortly.") + raise HTTPException( + status_code=503, detail="Graph is still initializing. Please retry shortly." + ) llm_response, context_sources, tools = parse_agent_output(output) @@ -410,7 +421,7 @@ async def get_response_stream(user_input: UserInput, db: Session | None) -> Any: conversation_uuid = uuid4() if conversation_uuid not in chat_history: - chat_history[conversation_uuid] = [] + add_conversation_to_history(conversation_uuid) inputs = { "messages": [ @@ -447,7 +458,7 @@ async def get_response_stream(user_input: UserInput, db: Session | None) -> Any: if msg: chunks.append(str(msg)) - yield str(msg) + "\n\n" + yield str(msg) + "\n\n" else: yield "Error: Graph is still initializing. Please retry shortly.\n\n" return diff --git a/backend/tests/test_api_conversations_streaming.py b/backend/tests/test_api_conversations_streaming.py index 193a250e..a3ecac91 100644 --- a/backend/tests/test_api_conversations_streaming.py +++ b/backend/tests/test_api_conversations_streaming.py @@ -507,3 +507,64 @@ async def mock_astream_events(*args, **kwargs): assert any("Sources:" in c for c in chunks) sources_chunk = [c for c in chunks if "Sources:" in c][0] assert sources_chunk.strip() == "Sources:" + + @pytest.mark.asyncio + async def test_get_response_stream_never_yields_none( + self, db_session: Session, mock_retriever_graph, sample_user_input: UserInput + ): + """Regression test: non-AIMessageChunk events must not leak the + literal string "None" into the client stream (issue #243, bug 2).""" + from src.api.routers.conversations import get_response_stream + + async def mock_astream_events(*args, **kwargs): + yield {"event": "on_chat_model_end", "data": {}} + # Non-AIMessageChunk event on the second LLM call: msg becomes None + yield { + "event": "on_chat_model_stream", + "data": {"chunk": "plain_string"}, + } + yield { + "event": "on_chat_model_stream", + "data": {"chunk": AIMessageChunk(content="Valid content")}, + } + + mock_retriever_graph.astream_events = mock_astream_events + + chunks = [] + async for chunk in get_response_stream(sample_user_input, db_session): + chunks.append(chunk) + + assert not any(c.strip() == "None" for c in chunks) + + +class TestInMemoryChatHistoryBounding: + """Regression tests: the in-memory chat_history store used when USE_DB is + false must be bounded, not grow without limit (issue #243, bug 4).""" + + @pytest.fixture(autouse=True) + def _clean_chat_history(self): + conversations.chat_history.clear() + yield + conversations.chat_history.clear() + + def test_add_conversation_to_history_evicts_oldest_when_over_capacity(self): + with patch.object(conversations, "MAX_IN_MEMORY_CONVERSATIONS", 3): + uuids = [uuid4() for _ in range(5)] + for u in uuids: + conversations.add_conversation_to_history(u) + + assert len(conversations.chat_history) == 3 + # The two oldest conversations should have been evicted + assert uuids[0] not in conversations.chat_history + assert uuids[1] not in conversations.chat_history + # The most recent ones remain + assert uuids[2] in conversations.chat_history + assert uuids[3] in conversations.chat_history + assert uuids[4] in conversations.chat_history + + def test_add_conversation_to_history_stays_bounded_over_many_inserts(self): + with patch.object(conversations, "MAX_IN_MEMORY_CONVERSATIONS", 10): + for _ in range(1000): + conversations.add_conversation_to_history(uuid4()) + + assert len(conversations.chat_history) == 10 diff --git a/backend/tests/test_process_html.py b/backend/tests/test_process_html.py index a4d63d7d..1397bea2 100644 --- a/backend/tests/test_process_html.py +++ b/backend/tests/test_process_html.py @@ -132,9 +132,7 @@ def test_process_html_split_without_chunk_size_raises_error(self): "builtins.open", mock_open(read_data='{"test.html": "https://example.com"}'), ): - with patch( - "src.tools.process_html.BSHTMLLoader" - ) as mock_loader: + with patch("src.tools.process_html.BSHTMLLoader") as mock_loader: mock_doc = Mock() mock_doc.metadata = {"source": "test.html"} mock_loader_instance = Mock() @@ -280,9 +278,7 @@ def test_process_html_real_file_structure(self): "builtins.open", mock_open(read_data='{"docs/html/test.html": "https://example.com"}'), ): - with patch( - "src.tools.process_html.BSHTMLLoader" - ) as mock_loader: + with patch("src.tools.process_html.BSHTMLLoader") as mock_loader: mock_doc = Mock() mock_doc.metadata = {"source": "test.html"} mock_doc.page_content = "Test Content" diff --git a/backend/tests/test_retriever_graph_simplified.py b/backend/tests/test_retriever_graph_simplified.py index d3316555..9fef5cea 100644 --- a/backend/tests/test_retriever_graph_simplified.py +++ b/backend/tests/test_retriever_graph_simplified.py @@ -1,6 +1,9 @@ import pytest from unittest.mock import Mock, patch +from langchain_core.messages import AIMessage +from langchain_core.runnables import RunnableLambda + from src.agents.retriever_graph import RetrieverGraph from src.agents.retriever_rag import ToolNode @@ -208,3 +211,86 @@ def test_route_with_empty_tools(self, mock_base_chain, mock_retriever_tools): result = graph.rag_route(state) assert result == "retrieve_general" + + +class TestRagAgentToolSelection: + """Test suite for the non-inbuilt-tool-calling branch of RAG.rag_agent.""" + + def _make_graph(self, fake_llm: RunnableLambda) -> RetrieverGraph: + with ( + patch("src.agents.retriever_rag.RetrieverTools") as mock_retriever_tools, + patch("src.agents.retriever_graph.BaseChain") as mock_base_chain, + ): + mock_chain_instance = Mock() + mock_base_chain.return_value = mock_chain_instance + mock_chain_instance.get_llm_chain.return_value = Mock() + + mock_tools_instance = Mock() + mock_retriever_tools.return_value = mock_tools_instance + mock_tools_instance.retrieve_cmds = Mock() + mock_tools_instance.retrieve_install = Mock() + mock_tools_instance.retrieve_general = Mock() + mock_tools_instance.retrieve_klayout_docs = Mock() + mock_tools_instance.retrieve_errinfo = Mock() + mock_tools_instance.retrieve_yosys_rtdocs = Mock() + + return RetrieverGraph( + llm_model=fake_llm, + embeddings_config={"type": "HF", "name": "test-model"}, + reranking_model_name="test-reranker", + inbuilt_tool_calling=False, + ) + + def _make_state(self) -> dict: + mock_message = Mock() + mock_message.content = "test query" + return { + "messages": [mock_message], + "context": [], + "context_list": [], + "tools": [], + "sources": [], + "urls": [], + "chat_history": "", + } + + def test_rag_agent_missing_tool_names_does_not_crash(self): + """Regression test: a response without "tool_names" must not raise + UnboundLocalError on `tool_calls` (issue #243, bug 1).""" + fake_llm = RunnableLambda(lambda _: AIMessage(content='{"foo": "bar"}')) + graph = self._make_graph(fake_llm) + + result = graph.rag_agent(self._make_state()) + + assert result == {"tools": []} + + def test_rag_agent_filters_all_invalid_tool_names(self): + """Regression test: filtering invalid tool names must not mutate the + list while iterating over it, which previously skipped entries + (issue #243, bug 1).""" + fake_llm = RunnableLambda( + lambda _: AIMessage( + content='{"tool_names": ["invalid1", "invalid2", "retrieve_cmds"]}' + ) + ) + graph = self._make_graph(fake_llm) + + result = graph.rag_agent(self._make_state()) + + assert result == {"tools": ["retrieve_cmds"]} + + def test_tool_descriptions_have_type_signature_stripped(self): + """Regression test: str.replace() result must be assigned back so the + noisy type signature is actually removed (issue #243, bug 3).""" + fake_llm = RunnableLambda(lambda _: AIMessage(content="{}")) + + with patch( + "src.agents.retriever_rag.render_text_description", + return_value="retrieve_cmds(query: str) -> Tuple[str, list[str], list[str]] - desc", + ): + graph = self._make_graph(fake_llm) + + assert ( + "(query: str) -> Tuple[str, list[str], list[str]]" + not in graph.tool_descriptions + )