From e38d006aa718d7167982346afdec58f7b03a63fb Mon Sep 17 00:00:00 2001 From: georgeatparallel Date: Tue, 20 Oct 2026 03:42:38 +0000 Subject: [PATCH 1/3] Add keyless Parallel search tools for unconfigured agents --- README.md | 4 + .../conversational_agent.py | 3 +- .../openai_functions_agent.py | 3 +- autochain/chain/base_chain.py | 4 +- autochain/tools/base.py | 2 +- autochain/tools/parallel_search.py | 182 ++++++++++ docs/tool.md | 38 +- tests/tools/test_parallel_search.py | 330 ++++++++++++++++++ 8 files changed, 561 insertions(+), 5 deletions(-) create mode 100644 autochain/tools/parallel_search.py create mode 100644 tests/tools/test_parallel_search.py diff --git a/README.md b/README.md index e8bb44e..446a6c1 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,10 @@ chain = Chain(agent=agent, memory=memory) print(chain.run("Write me a poem about AI")['message']) ``` +Omitting `tools` enables keyless Parallel web search and page fetching. Pass +`tools=[]` to disable tools, or supply your own list to keep full control of web +access. See [tool setup and privacy](./docs/tool.md#parallel-web-search-and-page-fetching). + Just like in LangChain, you can add a list of tools to the agent ```python diff --git a/autochain/agent/conversational_agent/conversational_agent.py b/autochain/agent/conversational_agent/conversational_agent.py index cf93324..ed6df30 100644 --- a/autochain/agent/conversational_agent/conversational_agent.py +++ b/autochain/agent/conversational_agent/conversational_agent.py @@ -19,6 +19,7 @@ from autochain.agent.structs import AgentAction, AgentFinish from autochain.models.base import BaseLanguageModel, Generation from autochain.tools.base import Tool +from autochain.tools.parallel_search import parallel_search_tools from autochain.utils import print_with_color logger = logging.getLogger(__name__) @@ -51,7 +52,7 @@ def from_llm_and_tools( **kwargs: Any, ) -> ConversationalAgent: """Construct an agent from an LLM and tools.""" - tools = tools or [] + tools = parallel_search_tools() if tools is None else tools template = cls.get_prompt_template( template=prompt_template, diff --git a/autochain/agent/openai_functions_agent/openai_functions_agent.py b/autochain/agent/openai_functions_agent/openai_functions_agent.py index 4faf99c..d83dd74 100644 --- a/autochain/agent/openai_functions_agent/openai_functions_agent.py +++ b/autochain/agent/openai_functions_agent/openai_functions_agent.py @@ -13,6 +13,7 @@ from autochain.agent.structs import AgentAction, AgentFinish from autochain.models.base import BaseLanguageModel, Generation from autochain.tools.base import Tool +from autochain.tools.parallel_search import parallel_search_tools from autochain.utils import print_with_color from colorama import Fore @@ -43,7 +44,7 @@ def from_llm_and_tools( min_confidence: int = 3, **kwargs: Any, ) -> OpenAIFunctionsAgent: - tools = tools or [] + tools = parallel_search_tools() if tools is None else tools allowed_tools = {tool.name: tool for tool in tools} _output_parser = output_parser or OpenAIFunctionOutputParser() diff --git a/autochain/chain/base_chain.py b/autochain/chain/base_chain.py index 7728e9c..cd3169b 100644 --- a/autochain/chain/base_chain.py +++ b/autochain/chain/base_chain.py @@ -13,6 +13,7 @@ from autochain.chain import constants from autochain.memory.base import BaseMemory from autochain.tools.base import Tool +from autochain.tools.parallel_search import search_conversation from pydantic import BaseModel logger = logging.getLogger(__name__) @@ -90,7 +91,8 @@ def run( inputs = self.prep_inputs(user_query) logger.info(f"\n Input to agent: {inputs}") try: - output = self._run(inputs) + with search_conversation(self.memory): + output = self._run(inputs) except (KeyboardInterrupt, Exception) as e: raise e diff --git a/autochain/tools/base.py b/autochain/tools/base.py index 629ddfc..ea6c5c7 100644 --- a/autochain/tools/base.py +++ b/autochain/tools/base.py @@ -102,7 +102,7 @@ def run( try: tool_args, tool_kwargs = self._to_args_and_kwargs(parsed_input) tool_output = self._run(*tool_args, **tool_kwargs) - except (Exception, KeyboardInterrupt) as e: + except Exception as e: raise ToolRunningError( message=f"Failed to run tool {self.name} due to {e}" ) from e diff --git a/autochain/tools/parallel_search.py b/autochain/tools/parallel_search.py new file mode 100644 index 0000000..d898935 --- /dev/null +++ b/autochain/tools/parallel_search.py @@ -0,0 +1,182 @@ +"""Keyless search and page extraction through Parallel's HTTP MCP endpoint.""" +import json +import socket +from contextlib import contextmanager +from contextvars import ContextVar +from typing import List +from urllib.error import HTTPError, URLError +from urllib.parse import urlsplit +from urllib.request import Request, urlopen +from uuid import uuid4 + +from autochain.tools.base import Tool + +ENDPOINT = "https://search.parallel.ai/mcp" +USER_AGENT = "AutoChain (https://github.com/Forethought-Technologies/AutoChain)" +_conversation = ContextVar("autochain_parallel_conversation", default=None) +_MEMORY_KEY = "_autochain_parallel_conversation" +_MAX_RESPONSE = 2 * 1024 * 1024 + + +@contextmanager +def search_conversation(memory): + """Bind search identity to chain memory; clearing memory starts a new chat.""" + session_id = memory.load_memory(_MEMORY_KEY) if memory is not None else None + if not session_id: + session_id = str(uuid4()) + if memory is not None: + memory.save_memory(_MEMORY_KEY, session_id) + token = _conversation.set(session_id) + try: + yield + finally: + _conversation.reset(token) + + +class _Client: + """Small synchronous Streamable HTTP client for this fixed server only.""" + + def __init__(self, timeout): + self.timeout = timeout + self.session = None + + def request(self, method, params, request_id=None): + message = {"jsonrpc": "2.0", "method": method, "params": params} + if request_id is not None: + message["id"] = request_id + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + "User-Agent": USER_AGENT, + } + if self.session: + headers["Mcp-Session-Id"] = self.session + headers["MCP-Protocol-Version"] = "2025-03-26" + request = Request(ENDPOINT, data=json.dumps(message).encode(), headers=headers) + with urlopen(request, timeout=self.timeout) as response: + if method == "initialize": + self.session = response.headers.get("Mcp-Session-Id") + if request_id is None: + return None + body = response.read(_MAX_RESPONSE + 1) + if len(body) > _MAX_RESPONSE: + raise ValueError("Parallel response exceeds 2 MiB") + if "text/event-stream" in response.headers.get("Content-Type", ""): + # Each SSE event carries one JSON-RPC message; ignore notifications. + messages = [] + for event in body.decode().replace("\r\n", "\n").split("\n\n"): + data = "\n".join( + line[5:].lstrip(" ") + for line in event.splitlines() + if line.startswith("data:") + ) + if data: + messages.append(json.loads(data)) + else: + messages = [json.loads(body)] + result = next((m for m in messages if m.get("id") == request_id), None) + if not result or result.get("jsonrpc") != "2.0": + raise ValueError("Invalid Parallel MCP response") + if "error" in result: + raise ValueError("Parallel MCP returned an error") + return result["result"] + + def call(self, name, arguments): + try: + initialized = self.request( + "initialize", + { + "protocolVersion": "2025-03-26", + "capabilities": {}, + "clientInfo": {"name": "autochain", "version": "0.0.5"}, + }, + 1, + ) + if initialized.get("protocolVersion") != "2025-03-26": + raise ValueError("Unsupported Parallel MCP protocol version") + self.request("notifications/initialized", {}) + tools = self.request("tools/list", {}, 2)["tools"] + if name not in [tool["name"] for tool in tools]: + raise ValueError("Parallel MCP tool is unavailable") + # Conversation identity is separate from the MCP transport session. + arguments["session_id"] = _conversation.get() or str(uuid4()) + result = self.request( + "tools/call", {"name": name, "arguments": arguments}, 3 + ) + text = "\n".join( + item["text"] + for item in result.get("content", []) + if item["type"] == "text" + ) + if result.get("isError"): + return "Parallel tool error: " + text + if not text: + raise ValueError("Parallel returned no text") + return text + except HTTPError as error: + return f"Parallel unavailable (HTTP {error.code}); retry later or use another configured tool." + except (URLError, TimeoutError, socket.timeout): + return "Parallel request failed or timed out; retry later." + except (ValueError, KeyError, TypeError): + return "Parallel returned an invalid MCP response; retry later." + finally: + if self.session: + request = Request( + ENDPOINT, + method="DELETE", + headers={ + "Mcp-Session-Id": self.session, + "MCP-Protocol-Version": "2025-03-26", + "User-Agent": USER_AGENT, + }, + ) + try: + with urlopen(request, timeout=self.timeout): + pass + except (OSError, ValueError): + pass # Session cleanup must not hide the tool's result. + + +def parallel_search_tools(timeout: float = 30.0) -> List[Tool]: + """Create search/fetch tools without connecting or reading ambient API keys. + + The timeout is in seconds per HTTP request. Standalone calls each start a + research run; calls dispatched by a Chain share its conversation identity. + """ + if timeout <= 0: + raise ValueError("timeout must be positive") + + # Plain functions keep compatibility with ChatOpenAI's signature inspection. + def web_search(query: str): + if not query.strip(): + return "Provide a nonempty search query." + return _Client(timeout).call( + "web_search", {"objective": query, "search_queries": [query]} + ) + + def web_fetch(url: str): + try: + parsed = urlsplit(url) + except ValueError: + return "Provide an HTTP or HTTPS URL without credentials." + if ( + parsed.scheme not in ("http", "https") + or not parsed.hostname + or parsed.username + or parsed.password + ): + return "Provide an HTTP or HTTPS URL without credentials." + return _Client(timeout).call("web_fetch", {"urls": [url]}) + + return [ + Tool( + name="web_search", + description="Search the web for current information. Input: a search query. Returns source URLs and excerpts. Powered by Parallel Search MCP.", + func=web_search, + ), + Tool( + name="web_fetch", + description="Read a web page for evidence. Input: an HTTP or HTTPS URL, usually from web_search. Returns page excerpts.", + func=web_fetch, + ), + ] diff --git a/docs/tool.md b/docs/tool.md index 52f35fd..1832c63 100644 --- a/docs/tool.md +++ b/docs/tool.md @@ -31,6 +31,42 @@ formatted into the prompt when using `OpenAIFunctionsAgent`. ## Tools included +### Parallel web search and page fetching + +`ConversationalAgent.from_llm_and_tools(llm=llm)` and +`OpenAIFunctionsAgent.from_llm_and_tools(llm=llm)` now include `web_search` and +`web_fetch` when `tools` is omitted or `None`. These tools use +[Parallel Search MCP](https://docs.parallel.ai/integrations/mcp/search-mcp) +at `https://search.parallel.ai/mcp`. Search runs in Fast mode with no account or +API key required. Anonymous use has service limits; failures are returned to the +agent as tool output without switching providers. + +The existing `ChatOpenAI` adapter uses OpenAI's pre-1.0 Python API. Use the +version recorded in `poetry.lock` (`openai==0.27.8`) for that adapter; the web +tools themselves do not require OpenAI or Parallel credentials. + +Pass `tools=[]` to disable all tools and keep the agent offline apart from model +inference. Any explicit tool list remains authoritative, including Google search +or custom tools. Constructing an agent does not make network requests. Calling +these tools sends the query or URL to Parallel; only enable them when web access +is appropriate for your application. + +To combine them with your own tools or adjust the per-request timeout: + +```python +from autochain.tools.parallel_search import parallel_search_tools + +tools = parallel_search_tools(timeout=20) + your_other_tools +agent = OpenAIFunctionsAgent.from_llm_and_tools(llm=llm, tools=tools) +``` + +The synchronous client uses Python's standard library, including on Python 3.8. +Each HTTP request has a 30-second timeout by default. A chain keeps one opaque +search conversation identifier in its memory across turns; clear the memory to +start a new conversation. Separate chains should use separate memory instances. +Standalone `Tool.run` calls each start a separate research run. Ctrl-C cancels +execution rather than converting the interruption to a tool error. + ### GoogleSearchTool Migrated from LangChain, which is also an example for user to easily migrate any tool from LangChain if needed. @@ -52,4 +88,4 @@ long term memory for the agent. ### LanceDBTool Internal search tool that can be used for long term memory of the agent or looking up relevant information that does not exists from the Internet. Currently, AutoChain supports `ChromaDB` as -long term memory for the agent. LanceDBTool is serverless, and does not require any setup. \ No newline at end of file +long term memory for the agent. LanceDBTool is serverless, and does not require any setup. diff --git a/tests/tools/test_parallel_search.py b/tests/tools/test_parallel_search.py new file mode 100644 index 0000000..60e666b --- /dev/null +++ b/tests/tools/test_parallel_search.py @@ -0,0 +1,330 @@ +import json +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from io import BytesIO +from unittest import mock +from urllib.error import HTTPError, URLError + +import pytest +from autochain.agent.conversational_agent.conversational_agent import ( + ConversationalAgent, +) +from autochain.agent.message import AIMessage, FunctionMessage +from autochain.agent.openai_functions_agent.openai_functions_agent import ( + OpenAIFunctionsAgent, +) +from autochain.chain.chain import Chain +from autochain.memory.buffer_memory import BufferMemory +from autochain.models.base import BaseLanguageModel, Generation, LLMResult +from autochain.models.chat_openai import convert_tool_to_dict +from autochain.tools.base import Tool +from autochain.tools.parallel_search import ENDPOINT, USER_AGENT, parallel_search_tools +from pydantic import PrivateAttr + + +class ScriptedModel(BaseLanguageModel): + """Deterministic model drives real function dispatch and consumes tool results.""" + + _step: int = PrivateAttr(default=0) + first_query: str = "Python docs" + + def generate(self, messages, functions=None, stop=None): + if len(messages) == 1 and "confidence" in messages[0].content: + message = AIMessage(content="5") + else: + errors = [ + m.content + for m in messages[-1:] + if isinstance(m, FunctionMessage) and m.content.startswith("Parallel ") + ] + if errors: + return LLMResult( + generations=[Generation(message=AIMessage(content=errors[-1]))] + ) + calls = [ + ("web_search", {"query": self.first_query}), + ("web_fetch", {"url": "https://docs.python.org/3/"}), + ("web_search", {"query": "Python releases"}), + ] + if self._step < len(calls): + name, arguments = calls[self._step] + assert name in [tool.name for tool in functions] + message = AIMessage( + content="", + function_call={"name": name, "arguments": json.dumps(arguments)}, + ) + else: + outputs = [ + m.content for m in messages if isinstance(m, FunctionMessage) + ] + assert any("https://docs.python.org/3/" in text for text in outputs) + assert any("Python documentation" in text for text in outputs) + message = AIMessage( + content="Python documentation: https://docs.python.org/3/" + ) + self._step += 1 + return LLMResult(generations=[Generation(message=message)]) + + +@contextmanager +def mcp_server(failure=None, sse=False): + """Observe actual request bytes below Tool.run, not a mocked client helper.""" + captured = [] + + def respond(request, timeout): + assert request.full_url == ENDPOINT + assert request.get_header("User-agent") == USER_AGENT + assert timeout == 30.0 + assert request.get_header("Authorization") is None + if request.get_method() == "DELETE": + return BytesIO() + body = json.loads(request.data) + captured.append(body) + method = body["method"] + if method != "initialize": + assert request.get_header("Mcp-session-id") == "transport-session" + assert request.get_header("Mcp-protocol-version") == "2025-03-26" + if method == "initialize": + result = {"protocolVersion": "2025-03-26"} + elif method == "notifications/initialized": + return BytesIO() + elif method == "tools/list": + result = {"tools": [{"name": "web_search"}, {"name": "web_fetch"}]} + else: + if failure: + raise failure + result = { + "content": [ + { + "type": "text", + "text": "Python documentation https://docs.python.org/3/", + } + ] + } + payload = json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": result}) + if sse: + payload = ( + 'event: message\ndata: {"jsonrpc":"2.0","method":"progress"}\n\n' + + "data: " + + payload + + "\n\n" + ) + response = BytesIO(payload.encode()) + response.headers = { + "Mcp-Session-Id": "transport-session", + "Content-Type": "text/event-stream" if sse else "application/json", + } + return response + + with mock.patch("autochain.tools.parallel_search.urlopen", side_effect=respond): + yield captured + + +def make_chain(tools=None): + agent = OpenAIFunctionsAgent.from_llm_and_tools(llm=ScriptedModel(), tools=tools) + return Chain(agent=agent, memory=BufferMemory()) + + +def tool_calls(captured): + return [m["params"] for m in captured if m["method"] == "tools/call"] + + +@pytest.mark.parametrize("factory", [ConversationalAgent, OpenAIFunctionsAgent]) +def test_defaults_explicit_choices_and_function_schemas(factory, monkeypatch): + for key in [ + "OPENAI_API_KEY", + "PARALLEL_API_KEY", + "GOOGLE_API_KEY", + "GOOGLE_CSE_ID", + ]: + monkeypatch.delenv(key, raising=False) + with mock.patch("autochain.tools.parallel_search.urlopen") as network: + agent = factory.from_llm_and_tools(llm=ScriptedModel()) + assert [t.name for t in agent.tools] == ["web_search", "web_fetch"] + assert factory.from_llm_and_tools(llm=ScriptedModel(), tools=[]).tools == [] + chosen = Tool( + name="my_search", description="Explicit provider", func=lambda query: query + ) + assert [ + t.name + for t in factory.from_llm_and_tools( + llm=ScriptedModel(), tools=[chosen] + ).tools + ] == ["my_search"] + network.assert_not_called() + for index, tool in enumerate(agent.tools): + argument = ["query", "url"][index] + spec = convert_tool_to_dict(tool) + assert spec["parameters"]["properties"] == {argument: {"type": "string"}} + assert spec["parameters"]["required"] == [argument] + + +@pytest.mark.parametrize("sse", [False, True]) +def test_agent_loop_consumes_results_and_preserves_conversation(sse): + chain = make_chain() + with mcp_server(sse=sse) as captured: + assert "https://docs.python.org/3/" in chain.run("Research Python")["message"] + calls = tool_calls(captured) + assert [c["name"] for c in calls] == ["web_search", "web_fetch", "web_search"] + assert calls[0]["arguments"]["search_queries"] == ["Python docs"] + assert calls[1]["arguments"]["urls"] == ["https://docs.python.org/3/"] + session = calls[0]["arguments"]["session_id"] + assert all(c["arguments"]["session_id"] == session for c in calls) + # A follow-up turn reuses the identity even with a fresh model instance. + chain.agent.llm = ScriptedModel() + chain.run("Follow up") + assert tool_calls(captured)[-1]["arguments"]["session_id"] == session + chain.memory.clear() + chain.agent.llm = ScriptedModel() + chain.run("New chat") + assert tool_calls(captured)[-1]["arguments"]["session_id"] != session + + +def test_interleaved_chains_share_tools_but_not_conversation_identity(): + tools = parallel_search_tools() + chains = [make_chain(tools), make_chain(tools)] + with mcp_server() as captured: + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(lambda c: c.run("Research Python"), chains)) + assert all("Python documentation" in r["message"] for r in results) + ids = [c["arguments"]["session_id"] for c in tool_calls(captured)] + assert len(set(ids)) == 2 + assert all(ids.count(session) == 3 for session in set(ids)) + + +@pytest.mark.parametrize( + "failure,expected", + [ + (HTTPError(ENDPOINT, 429, "Quota", {}, None), "HTTP 429"), + (HTTPError(ENDPOINT, 401, "Unauthorized", {}, None), "HTTP 401"), + (URLError("offline"), "failed or timed out"), + (TimeoutError(), "failed or timed out"), + ], +) +def test_agent_consumes_errors_without_provider_fallback(failure, expected): + chain = make_chain() + with mcp_server(failure=failure) as captured: + assert expected in chain.run("Research Python")["message"] + assert len(tool_calls(captured)) == 1 + + +@pytest.mark.parametrize( + "failure", [TimeoutError(), HTTPError(ENDPOINT, 429, "Quota", {}, None)] +) +def test_retry_retains_conversation_and_cancel_propagates(failure): + chain = make_chain() + with mcp_server(failure=failure) as captured: + # A single iteration retains the failed request as an observation. + chain.max_iterations = 1 + chain.run("Research Python") + session = tool_calls(captured)[0]["arguments"]["session_id"] + chain.agent.llm = ScriptedModel(first_query="Python documentation retries") + with mcp_server() as captured: + chain.run("Retry") + assert tool_calls(captured)[0]["arguments"]["session_id"] == session + chain.agent.llm = ScriptedModel(first_query="Cancel Python docs") + with mcp_server(failure=KeyboardInterrupt()): + with pytest.raises(KeyboardInterrupt): + chain.run("Cancel") + + +@pytest.mark.parametrize( + "url", + [ + "file:///etc/passwd", + "ftp://example.com", + "https://user:secret@example.com", + "not a URL", + "https://[invalid", + ], +) +def test_invalid_urls_never_connect(url): + with mock.patch("autochain.tools.parallel_search.urlopen") as network: + assert "HTTP or HTTPS" in parallel_search_tools()[1].run({"url": url}) + network.assert_not_called() + + +class OneToolModel(BaseLanguageModel): + action_name: str = "web_search" + + def generate(self, messages, functions=None, stop=None): + if len(messages) == 1 and "confidence" in messages[0].content: + message = AIMessage(content="5") + elif isinstance(messages[-1], FunctionMessage): + message = AIMessage(content=messages[-1].content) + else: + message = AIMessage( + content="", + function_call={ + "name": self.action_name, + "arguments": '{"query": "Python docs"}', + }, + ) + return LLMResult(generations=[Generation(message=message)]) + + +def test_disabled_tools_do_not_dispatch(): + agent = OpenAIFunctionsAgent.from_llm_and_tools(llm=OneToolModel(), tools=[]) + chain = Chain(agent=agent, memory=BufferMemory()) + with mock.patch("autochain.tools.parallel_search.urlopen") as network: + assert "not supported" in chain.run("Search anyway")["message"] + network.assert_not_called() + + +def test_explicit_provider_is_used_without_parallel_fallback(): + def configured_search(query): + return "Configured provider's result or failure" + + selected = Tool( + name="web_search", description="Saved provider", func=configured_search + ) + agent = OpenAIFunctionsAgent.from_llm_and_tools( + llm=OneToolModel(), tools=[selected] + ) + with mock.patch("autochain.tools.parallel_search.urlopen") as network: + result = Chain(agent=agent, memory=BufferMemory()).run("Search") + assert result["message"] == "Configured provider's result or failure" + network.assert_not_called() + + +def test_malformed_response_and_timeout_setting(): + with pytest.raises(ValueError): + parallel_search_tools(timeout=0) + with mock.patch("autochain.tools.parallel_search.urlopen") as network: + response = BytesIO(b"not JSON") + response.headers = {} + network.return_value = response + assert "invalid MCP response" in parallel_search_tools()[0].run("Python docs") + + +class ScriptedConversationModel(BaseLanguageModel): + def generate(self, messages, functions=None, stop=None): + prompt = messages[0].content + if "has assistant finish helping" in prompt: + text = "no" + elif "has_arg_value" in prompt: + text = json.dumps({"has_arg_value": "Yes", "clarifying_question": ""}) + elif "Python documentation https://docs.python.org/3/" in prompt: + text = json.dumps( + { + "thoughts": {"need_use_tool": "No"}, + "response": "Python documentation https://docs.python.org/3/", + } + ) + else: + text = json.dumps( + { + "thoughts": {"need_use_tool": "Yes"}, + "tool": {"name": "web_search", "args": {"query": "Python docs"}}, + } + ) + return LLMResult(generations=[Generation(message=AIMessage(content=text))]) + + +def test_conversational_agent_consumes_search_results(): + agent = ConversationalAgent.from_llm_and_tools(llm=ScriptedConversationModel()) + chain = Chain(agent=agent, memory=BufferMemory()) + with mcp_server() as captured: + result = chain.run("Find Python documentation") + assert "Python documentation https://docs.python.org/3/" in result["message"] + assert len(tool_calls(captured)) == 1 From 01de98c36556c7b1da6b8deae60c5697b41c92d6 Mon Sep 17 00:00:00 2001 From: georgeatparallel Date: Tue, 20 Oct 2026 03:42:55 +0000 Subject: [PATCH 2/3] Name HTTP error fixtures for Python 3.8 pytest compatibility --- tests/tools/test_parallel_search.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/tools/test_parallel_search.py b/tests/tools/test_parallel_search.py index 60e666b..3f48cff 100644 --- a/tests/tools/test_parallel_search.py +++ b/tests/tools/test_parallel_search.py @@ -200,6 +200,7 @@ def test_interleaved_chains_share_tools_but_not_conversation_identity(): (URLError("offline"), "failed or timed out"), (TimeoutError(), "failed or timed out"), ], + ids=["quota", "authentication", "network", "timeout"], ) def test_agent_consumes_errors_without_provider_fallback(failure, expected): chain = make_chain() @@ -209,7 +210,9 @@ def test_agent_consumes_errors_without_provider_fallback(failure, expected): @pytest.mark.parametrize( - "failure", [TimeoutError(), HTTPError(ENDPOINT, 429, "Quota", {}, None)] + "failure", + [TimeoutError(), HTTPError(ENDPOINT, 429, "Quota", {}, None)], + ids=["timeout", "quota"], ) def test_retry_retains_conversation_and_cancel_propagates(failure): chain = make_chain() From a5aa0da1cd18aa3eb3eff69c83d40f7f11574649 Mon Sep 17 00:00:00 2001 From: georgeatparallel Date: Tue, 20 Oct 2026 03:46:07 +0000 Subject: [PATCH 3/3] Isolate conversation identity from retrieved memory and finish SSE at response --- autochain/memory/base.py | 14 +++- autochain/memory/buffer_memory.py | 1 + autochain/memory/long_term_memory.py | 1 + autochain/memory/redis_memory.py | 17 ++++ autochain/tools/parallel_search.py | 56 ++++++++----- docs/tool.md | 2 + tests/tools/test_parallel_search.py | 116 ++++++++++++++++++++++++++- 7 files changed, 183 insertions(+), 24 deletions(-) diff --git a/autochain/memory/base.py b/autochain/memory/base.py index c9f6259..4f59f7b 100644 --- a/autochain/memory/base.py +++ b/autochain/memory/base.py @@ -8,8 +8,9 @@ Optional, Union, ) +from uuid import UUID, uuid4 -from pydantic import BaseModel +from pydantic import BaseModel, Field from autochain.agent.message import ChatMessageHistory, MessageType @@ -17,6 +18,17 @@ class BaseMemory(BaseModel, ABC): """Base interface for memory in chains.""" + conversation_id: UUID = Field(default_factory=uuid4) + """Opaque conversation metadata, separate from retrievable user memories.""" + + def get_conversation_id(self) -> str: + """Return conversation identity without semantic retrieval.""" + return str(self.conversation_id) + + def reset_conversation_id(self) -> None: + """Start a fresh identity. Custom memory implementations call this in clear().""" + self.conversation_id = uuid4() + @abstractmethod def load_memory( self, key: Union[str, None] = None, default: Optional[Any] = None, **kwargs: Any diff --git a/autochain/memory/buffer_memory.py b/autochain/memory/buffer_memory.py index a8fb13f..bfb7315 100644 --- a/autochain/memory/buffer_memory.py +++ b/autochain/memory/buffer_memory.py @@ -36,5 +36,6 @@ def save_conversation( def clear(self) -> None: """Clear memory contents.""" + self.reset_conversation_id() self.conversation_history.clear() self.kv_memory = {} diff --git a/autochain/memory/long_term_memory.py b/autochain/memory/long_term_memory.py index 33493a7..68e6ace 100644 --- a/autochain/memory/long_term_memory.py +++ b/autochain/memory/long_term_memory.py @@ -66,6 +66,7 @@ def save_conversation( def clear(self) -> None: """Clear memory contents.""" + self.reset_conversation_id() self.conversation_history.clear() self.long_term_memory.clear_index() self.kv_memory = {} diff --git a/autochain/memory/redis_memory.py b/autochain/memory/redis_memory.py index de090e0..a9ec670 100644 --- a/autochain/memory/redis_memory.py +++ b/autochain/memory/redis_memory.py @@ -1,5 +1,6 @@ import pickle from typing import Any, Optional, Dict +from uuid import UUID from autochain.agent.message import ( ChatMessageHistory, @@ -51,6 +52,21 @@ def save_memory(self, key: str, value: Any) -> None: pickled = pickle.dumps(value) self.redis_client.set(key, pickled, ex=self.expire_time) + def get_conversation_id(self) -> str: + """Persist typed metadata independently of pickled conversation contents.""" + key = self.redis_key_prefix + ":_autochain_conversation_id" + saved = self.redis_client.get(key) + try: + if isinstance(saved, bytes): + saved = saved.decode("ascii") + if not isinstance(saved, str): + raise ValueError("Missing conversation identity") + self.conversation_id = UUID(saved) + except (ValueError, UnicodeError): + pass + self.redis_client.set(key, str(self.conversation_id), ex=self.expire_time) + return str(self.conversation_id) + def save_conversation( self, message: str, message_type: MessageType, **kwargs ) -> None: @@ -77,3 +93,4 @@ def clear(self) -> None: """Clear redis memory.""" for key in self.redis_client.keys(f"{self.redis_key_prefix}:*"): self.redis_client.delete(key) + self.reset_conversation_id() diff --git a/autochain/tools/parallel_search.py b/autochain/tools/parallel_search.py index d898935..477fd4d 100644 --- a/autochain/tools/parallel_search.py +++ b/autochain/tools/parallel_search.py @@ -14,18 +14,13 @@ ENDPOINT = "https://search.parallel.ai/mcp" USER_AGENT = "AutoChain (https://github.com/Forethought-Technologies/AutoChain)" _conversation = ContextVar("autochain_parallel_conversation", default=None) -_MEMORY_KEY = "_autochain_parallel_conversation" _MAX_RESPONSE = 2 * 1024 * 1024 @contextmanager def search_conversation(memory): """Bind search identity to chain memory; clearing memory starts a new chat.""" - session_id = memory.load_memory(_MEMORY_KEY) if memory is not None else None - if not session_id: - session_id = str(uuid4()) - if memory is not None: - memory.save_memory(_MEMORY_KEY, session_id) + session_id = memory.get_conversation_id() if memory is not None else str(uuid4()) token = _conversation.set(session_id) try: yield @@ -58,27 +53,44 @@ def request(self, method, params, request_id=None): self.session = response.headers.get("Mcp-Session-Id") if request_id is None: return None - body = response.read(_MAX_RESPONSE + 1) - if len(body) > _MAX_RESPONSE: - raise ValueError("Parallel response exceeds 2 MiB") if "text/event-stream" in response.headers.get("Content-Type", ""): - # Each SSE event carries one JSON-RPC message; ignore notifications. - messages = [] - for event in body.decode().replace("\r\n", "\n").split("\n\n"): - data = "\n".join( - line[5:].lstrip(" ") - for line in event.splitlines() - if line.startswith("data:") - ) - if data: - messages.append(json.loads(data)) + # Stop at the matching event even if the server keeps SSE open. + size, data = 0, [] + result = None + while size <= _MAX_RESPONSE: + line = response.readline(_MAX_RESPONSE - size + 1) + size += len(line) + if size > _MAX_RESPONSE: + raise ValueError("Parallel response exceeds 2 MiB") + if not line: + break + text = line.decode().rstrip("\r\n") + if text.startswith("data:"): + data.append(text[5:].lstrip(" ")) + elif not text and data: + message = json.loads("\n".join(data)) + data = [] + if ( + isinstance(message, dict) + and message.get("id") == request_id + ): + result = message + break else: - messages = [json.loads(body)] - result = next((m for m in messages if m.get("id") == request_id), None) - if not result or result.get("jsonrpc") != "2.0": + body = response.read(_MAX_RESPONSE + 1) + if len(body) > _MAX_RESPONSE: + raise ValueError("Parallel response exceeds 2 MiB") + result = json.loads(body) + if ( + not isinstance(result, dict) + or result.get("id") != request_id + or result.get("jsonrpc") != "2.0" + ): raise ValueError("Invalid Parallel MCP response") if "error" in result: raise ValueError("Parallel MCP returned an error") + if not isinstance(result.get("result"), dict): + raise ValueError("Invalid Parallel MCP result") return result["result"] def call(self, name, arguments): diff --git a/docs/tool.md b/docs/tool.md index 1832c63..419e54d 100644 --- a/docs/tool.md +++ b/docs/tool.md @@ -66,6 +66,8 @@ search conversation identifier in its memory across turns; clear the memory to start a new conversation. Separate chains should use separate memory instances. Standalone `Tool.run` calls each start a separate research run. Ctrl-C cancels execution rather than converting the interruption to a tool error. +Custom memory implementations should keep `get_conversation_id()` separate from +semantic retrieval and call `reset_conversation_id()` in their `clear()` method. ### GoogleSearchTool Migrated from LangChain, which is also an example for user to easily migrate any tool from diff --git a/tests/tools/test_parallel_search.py b/tests/tools/test_parallel_search.py index 3f48cff..99001df 100644 --- a/tests/tools/test_parallel_search.py +++ b/tests/tools/test_parallel_search.py @@ -1,9 +1,14 @@ +import importlib.util import json +import sys +import types from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from io import BytesIO +from pathlib import Path from unittest import mock from urllib.error import HTTPError, URLError +from uuid import UUID import pytest from autochain.agent.conversational_agent.conversational_agent import ( @@ -14,12 +19,16 @@ OpenAIFunctionsAgent, ) from autochain.chain.chain import Chain +from autochain.memory import base as memory_module from autochain.memory.buffer_memory import BufferMemory +from autochain.memory.redis_memory import RedisMemory from autochain.models.base import BaseLanguageModel, Generation, LLMResult from autochain.models.chat_openai import convert_tool_to_dict from autochain.tools.base import Tool +from autochain.tools.internal_search.base_search_tool import BaseSearchTool from autochain.tools.parallel_search import ENDPOINT, USER_AGENT, parallel_search_tools from pydantic import PrivateAttr +from redis import Redis class ScriptedModel(BaseLanguageModel): @@ -109,7 +118,17 @@ def respond(request, timeout): + payload + "\n\n" ) - response = BytesIO(payload.encode()) + + class OpenEventStream(BytesIO): + def readline(self, size=-1): + line = super().readline(size) + if not line: + raise TimeoutError("SSE connection remains open after the result") + return line + + response = ( + OpenEventStream(payload.encode()) if sse else BytesIO(payload.encode()) + ) response.headers = { "Mcp-Session-Id": "transport-session", "Content-Type": "text/event-stream" if sse else "application/json", @@ -331,3 +350,98 @@ def test_conversational_agent_consumes_search_results(): result = chain.run("Find Python documentation") assert "Python documentation https://docs.python.org/3/" in result["message"] assert len(tool_calls(captured)) == 1 + + +def test_long_term_memory_never_retrieves_conversation_metadata(): + # Load the actual memory implementation without unrelated optional DB SDKs. + backend_modules = {} + for module_name, search_name, doc_name in [ + ("chromadb_tool", "ChromaDBSearch", "ChromaDoc"), + ("pinecone_tool", "PineconeSearch", "PineconeDoc"), + ("lancedb_tool", "LanceDBSeach", "LanceDBDoc"), + ]: + name = "autochain.tools.internal_search." + module_name + module = types.ModuleType(name) + setattr(module, search_name, type(search_name, (BaseSearchTool,), {})) + setattr(module, doc_name, type(doc_name, (), {})) + backend_modules[name] = module + spec = importlib.util.spec_from_file_location( + "_tested_long_term_memory", + Path(memory_module.__file__).with_name("long_term_memory.py"), + ) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict(sys.modules, backend_modules): + spec.loader.exec_module(module) + + class Backend(BaseSearchTool): + _queries = PrivateAttr(default_factory=list) + + def _run(self, query, top_k=1): + return self.run({"query": query, "top_k": top_k}) + + def add_docs(self, docs, **kwargs): + pass + + def run(self, arguments): + self._queries.append(arguments["query"]) + if arguments["query"] == "intermediate_steps": + return [] + return "Doc 0: private customer information" + + def clear_index(self): + pass + + memory = module.LongTermMemory(long_term_memory=Backend()) + agent = OpenAIFunctionsAgent.from_llm_and_tools(llm=OneToolModel()) + chain = Chain(agent=agent, memory=memory) + with mcp_server() as captured: + assert "Python documentation" in chain.run("Research")["message"] + first = tool_calls(captured)[0]["arguments"]["session_id"] + assert str(UUID(first)) == first + assert ( + "_autochain_parallel_conversation" not in memory.long_term_memory._queries + ) + assert all( + not query.startswith("_autochain") + for query in memory.long_term_memory._queries + ) + memory.clear() + chain.run("New conversation") + assert tool_calls(captured)[-1]["arguments"]["session_id"] != first + + +def test_redis_conversations_persist_and_reset_without_sharing_identity(): + import fnmatch + + store = {} + client = mock.MagicMock(spec=Redis) + client.get.side_effect = store.get + + def put(key, value, **kwargs): + store[key] = value.encode() if isinstance(value, str) else value + + client.set.side_effect = put + client.keys.side_effect = lambda pattern: [ + k for k in store if fnmatch.fnmatch(k, pattern) + ] + client.delete.side_effect = lambda key: store.pop(key, None) + shared_tools = parallel_search_tools() + chains = [ + Chain( + agent=OpenAIFunctionsAgent.from_llm_and_tools( + llm=OneToolModel(), tools=shared_tools + ), + memory=RedisMemory(redis_key_prefix=prefix, redis_client=client), + ) + for prefix in ["conversation-one", "conversation-two"] + ] + with mcp_server() as captured: + for chain in chains: + chain.run("Research") + ids = [c["arguments"]["session_id"] for c in tool_calls(captured)] + assert len(set(ids)) == 2 + resumed = RedisMemory(redis_key_prefix="conversation-one", redis_client=client) + assert resumed.get_conversation_id() == ids[0] + chains[0].memory.clear() + chains[0].run("New chat") + assert tool_calls(captured)[-1]["arguments"]["session_id"] not in ids