From 3a1d5061c29f1c462498ba1fc403003c0b1b4dc3 Mon Sep 17 00:00:00 2001 From: Connor Shorten Date: Thu, 21 May 2026 08:46:40 -0400 Subject: [PATCH 1/2] first look --- test/query_agent/test_query_model.py | 94 ++++++++++++++++++++++++++++ weaviate_agents/query/query_agent.py | 19 ++++++ 2 files changed, 113 insertions(+) diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 3bb0c76..865c621 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -2119,3 +2119,97 @@ async def test_async_suggest_queries_failure(monkeypatch): str(exc_info.value) == "{'error': {'message': 'Test error message', 'code': 'test_error_code', 'details': {'info': 'test detail'}}}" ) + + +def test_suggest_queries_with_conversation(monkeypatch): + captured = {} + + def fake_post_with_capture(url, headers=None, json=None, timeout=None): + captured["json"] = json + return FakeResponse(200, FAKE_SUGGEST_QUERIES_SUCCESS_JSON) + + monkeypatch.setattr(httpx, "post", fake_post_with_capture) + dummy_client = DummyClient() + agent = QueryAgent( + dummy_client, ["test_collection"], agents_host="http://dummy-agent" + ) + agent._connection = dummy_client + agent._headers = dummy_client.additional_headers + + chat_messages: list[ChatMessage] = [ + {"role": "user", "content": "What topics are covered?"}, + {"role": "assistant", "content": "The collection covers ML and economics."}, + ] + + result = agent.suggest_queries(["test_collection"], conversation=chat_messages) + + assert isinstance(result, SuggestQueryResponse) + assert captured["json"]["conversation"] == {"messages": chat_messages} + + +def test_suggest_queries_without_conversation(monkeypatch): + captured = {} + + def fake_post_with_capture(url, headers=None, json=None, timeout=None): + captured["json"] = json + return FakeResponse(200, FAKE_SUGGEST_QUERIES_SUCCESS_JSON) + + monkeypatch.setattr(httpx, "post", fake_post_with_capture) + dummy_client = DummyClient() + agent = QueryAgent( + dummy_client, ["test_collection"], agents_host="http://dummy-agent" + ) + agent._connection = dummy_client + agent._headers = dummy_client.additional_headers + + agent.suggest_queries(["test_collection"]) + + assert "conversation" not in captured["json"] + + +async def test_async_suggest_queries_with_conversation(monkeypatch): + captured = {} + + async def fake_async_post_with_capture(*args, **kwargs): + captured["json"] = kwargs.get("json") + return FakeResponse(200, FAKE_SUGGEST_QUERIES_SUCCESS_JSON) + + monkeypatch.setattr(httpx.AsyncClient, "post", fake_async_post_with_capture) + dummy_client = DummyClient() + agent = AsyncQueryAgent( + dummy_client, ["test_collection"], agents_host="http://dummy-agent" + ) + agent._connection = dummy_client + agent._headers = dummy_client.additional_headers + + chat_messages: list[ChatMessage] = [ + {"role": "user", "content": "What topics are covered?"}, + {"role": "assistant", "content": "The collection covers ML and economics."}, + ] + + result = await agent.suggest_queries( + ["test_collection"], conversation=chat_messages + ) + + assert isinstance(result, SuggestQueryResponse) + assert captured["json"]["conversation"] == {"messages": chat_messages} + + +async def test_async_suggest_queries_without_conversation(monkeypatch): + captured = {} + + async def fake_async_post_with_capture(*args, **kwargs): + captured["json"] = kwargs.get("json") + return FakeResponse(200, FAKE_SUGGEST_QUERIES_SUCCESS_JSON) + + monkeypatch.setattr(httpx.AsyncClient, "post", fake_async_post_with_capture) + dummy_client = DummyClient() + agent = AsyncQueryAgent( + dummy_client, ["test_collection"], agents_host="http://dummy-agent" + ) + agent._connection = dummy_client + agent._headers = dummy_client.additional_headers + + await agent.suggest_queries(["test_collection"]) + + assert "conversation" not in captured["json"] diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index d3e9a64..ab9cff6 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -508,6 +508,7 @@ def suggest_queries( collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, num_queries: int = 3, instructions: Optional[str] = None, + conversation: Optional[list[ChatMessage]] = None, ) -> Union[SuggestQueryResponse, Coroutine[Any, Any, SuggestQueryResponse]]: pass @@ -1074,6 +1075,7 @@ def suggest_queries( collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, num_queries: int = 3, instructions: Optional[str] = None, + conversation: Optional[list[ChatMessage]] = None, ) -> SuggestQueryResponse: """Suggest queries for the data in your collections. @@ -1090,6 +1092,10 @@ def suggest_queries( instructions: Optional natural language guidance for the style, topic, or language of the suggested queries. This is supplied in addition to the agent's system instructions. + conversation: + Optional list of chat messages representing a prior conversation. + When provided, the suggested queries will be contextualised as + follow-up questions to the conversation. Returns: An instance of :class:`~weaviate_agents.query.classes.response.SuggestQueryResponse` which @@ -1136,6 +1142,10 @@ def suggest_queries( } if instructions is not None: request_body["instructions"] = instructions + if conversation is not None: + request_body["conversation"] = ConversationContext( + messages=conversation + ).model_dump(mode="json") response = httpx.post( self.query_url + "/suggest_queries", @@ -1714,6 +1724,7 @@ async def suggest_queries( collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, num_queries: int = 3, instructions: Optional[str] = None, + conversation: Optional[list[ChatMessage]] = None, ) -> SuggestQueryResponse: """Suggest queries for the data in your collections. @@ -1730,6 +1741,10 @@ async def suggest_queries( instructions: Optional natural language guidance for the style, topic, or language of the suggested queries. This is supplied in addition to the agent's system instructions. + conversation: + Optional list of chat messages representing a prior conversation. + When provided, the suggested queries will be contextualised as + follow-up questions to the conversation. Returns: An instance of :class:`~weaviate_agents.query.classes.response.SuggestQueryResponse` which @@ -1776,6 +1791,10 @@ async def suggest_queries( } if instructions is not None: request_body["instructions"] = instructions + if conversation is not None: + request_body["conversation"] = ConversationContext( + messages=conversation + ).model_dump(mode="json") async with httpx.AsyncClient() as client: response = await client.post( From 690143e798feb9c85f38889b6042abf2eb02554a Mon Sep 17 00:00:00 2001 From: Connor Shorten Date: Tue, 26 May 2026 07:32:04 -0400 Subject: [PATCH 2/2] fix conversation_context in request_body --- test/query_agent/test_query_model.py | 8 ++++---- weaviate_agents/query/query_agent.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 865c621..65b13a9 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -2144,7 +2144,7 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): result = agent.suggest_queries(["test_collection"], conversation=chat_messages) assert isinstance(result, SuggestQueryResponse) - assert captured["json"]["conversation"] == {"messages": chat_messages} + assert captured["json"]["conversation_context"] == {"messages": chat_messages} def test_suggest_queries_without_conversation(monkeypatch): @@ -2164,7 +2164,7 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): agent.suggest_queries(["test_collection"]) - assert "conversation" not in captured["json"] + assert "conversation_context" not in captured["json"] async def test_async_suggest_queries_with_conversation(monkeypatch): @@ -2192,7 +2192,7 @@ async def fake_async_post_with_capture(*args, **kwargs): ) assert isinstance(result, SuggestQueryResponse) - assert captured["json"]["conversation"] == {"messages": chat_messages} + assert captured["json"]["conversation_context"] == {"messages": chat_messages} async def test_async_suggest_queries_without_conversation(monkeypatch): @@ -2212,4 +2212,4 @@ async def fake_async_post_with_capture(*args, **kwargs): await agent.suggest_queries(["test_collection"]) - assert "conversation" not in captured["json"] + assert "conversation_context" not in captured["json"] diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index ab9cff6..5797c24 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -1143,7 +1143,7 @@ def suggest_queries( if instructions is not None: request_body["instructions"] = instructions if conversation is not None: - request_body["conversation"] = ConversationContext( + request_body["conversation_context"] = ConversationContext( messages=conversation ).model_dump(mode="json") @@ -1792,7 +1792,7 @@ async def suggest_queries( if instructions is not None: request_body["instructions"] = instructions if conversation is not None: - request_body["conversation"] = ConversationContext( + request_body["conversation_context"] = ConversationContext( messages=conversation ).model_dump(mode="json")