From 179dea16ff26583955f6bc4800813c614e0702ee Mon Sep 17 00:00:00 2001 From: Jack Luar Date: Sun, 12 Jul 2026 03:58:51 +0000 Subject: [PATCH] fix: resolve 4 runtime bugs from issue #243 salvages PR 258/308 Signed-off-by: Jack Luar Co-authored-by: sanjibani Co-authored-by: Amitesh Gupta --- backend/src/agents/retriever_rag.py | 12 +-- backend/src/api/routers/conversations.py | 21 +++-- .../tests/test_api_conversations_streaming.py | 56 +++++++++++++ .../tests/test_retriever_graph_simplified.py | 78 +++++++++++++++++++ 4 files changed, 157 insertions(+), 10 deletions(-) diff --git a/backend/src/agents/retriever_rag.py b/backend/src/agents/retriever_rag.py index a90260cf..9b2a595b 100644 --- a/backend/src/agents/retriever_rag.py +++ b/backend/src/agents/retriever_rag.py @@ -107,7 +107,7 @@ 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 +159,14 @@ def rag_agent(self, state: AgentState) -> dict[str, list[Any]]: ) return {"tools": []} + tool_calls: list[str] = [] 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: + raw_tool_calls = response.get("tool_names", []) # type: ignore + for tool in raw_tool_calls: + if tool in self.tool_names: + tool_calls.append(tool) + else: logging.warning(f"Tool {tool} not found in tool list.") - tool_calls.remove(tool) else: logging.warning(str(response)) logging.warning("Tool selection failed. Returning empty tool list.") 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..8bf3a761 100644 --- a/backend/tests/test_api_conversations_streaming.py +++ b/backend/tests/test_api_conversations_streaming.py @@ -507,3 +507,59 @@ 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: non-AIMessageChunk events must not leak "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": {}} + 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: in-memory chat_history must be bounded when USE_DB is false (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_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 + assert uuids[0] not in conversations.chat_history + assert uuids[1] not in conversations.chat_history + 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_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_retriever_graph_simplified.py b/backend/tests/test_retriever_graph_simplified.py index d3316555..92aeca6c 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,78 @@ def test_route_with_empty_tools(self, mock_base_chain, mock_retriever_tools): result = graph.rag_route(state) assert result == "retrieve_general" + + +class TestRagAgentToolSelection: + """Regression tests 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: response without "tool_names" must not raise UnboundLocalError (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_invalid_tool_names(self): + """Regression: invalid tools must not be included without skipping adjacent valid ones (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_rag_agent_all_invalid_tool_names_returns_empty(self): + """Regression: all-invalid tool list must return empty tools, not crash (issue #243, bug 1).""" + fake_llm = RunnableLambda( + lambda _: AIMessage(content='{"tool_names": ["bad1", "bad2"]}') + ) + graph = self._make_graph(fake_llm) + + result = graph.rag_agent(self._make_state()) + + assert result == {"tools": []}