Skip to content

Commit 394451a

Browse files
dannyjameswilliamsdannyjameswilliams
committed
feat: optional arg on .ask and .ask_stream methods
1 parent 73d46c2 commit 394451a

2 files changed

Lines changed: 37 additions & 5 deletions

File tree

‎uv.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎weaviate_agents/query/query_agent.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,7 @@ def ask(
183183
self,
184184
query: Union[str, list[ChatMessage]],
185185
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
186+
result_evaluation: Literal["llm", "none"] = "none",
186187
) -> Union[AskModeResponse, Coroutine[Any, Any, AskModeResponse]]:
187188
"""Run the Query Agent ask mode."""
188189
pass
@@ -281,6 +282,7 @@ def ask_stream(
281282
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
282283
include_progress: Literal[True] = True,
283284
include_final_state: Literal[True] = True,
285+
result_evaluation: Literal["llm", "none"] = "none",
284286
) -> Union[
285287
Generator[Union[ProgressMessage, StreamedTokens, AskModeResponse], None, None],
286288
AsyncGenerator[Union[ProgressMessage, StreamedTokens, AskModeResponse], None],
@@ -293,6 +295,7 @@ def ask_stream(
293295
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
294296
include_progress: Literal[True] = True,
295297
include_final_state: Literal[False] = False,
298+
result_evaluation: Literal["llm", "none"] = "none",
296299
) -> Union[
297300
Generator[Union[ProgressMessage, StreamedTokens], None, None],
298301
AsyncGenerator[Union[ProgressMessage, StreamedTokens], None],
@@ -305,6 +308,7 @@ def ask_stream(
305308
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
306309
include_progress: Literal[False] = False,
307310
include_final_state: Literal[True] = True,
311+
result_evaluation: Literal["llm", "none"] = "none",
308312
) -> Union[
309313
Generator[Union[StreamedTokens, AskModeResponse], None, None],
310314
AsyncGenerator[Union[StreamedTokens, AskModeResponse], None],
@@ -317,6 +321,7 @@ def ask_stream(
317321
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
318322
include_progress: Literal[False] = False,
319323
include_final_state: Literal[False] = False,
324+
result_evaluation: Literal["llm", "none"] = "none",
320325
) -> Union[
321326
Generator[StreamedTokens, None, None],
322327
AsyncGenerator[StreamedTokens, None],
@@ -329,6 +334,7 @@ def ask_stream(
329334
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
330335
include_progress: bool = True,
331336
include_final_state: bool = True,
337+
result_evaluation: Literal["llm", "none"] = "none",
332338
) -> Union[
333339
Generator[Union[ProgressMessage, StreamedTokens, AskModeResponse], None, None],
334340
AsyncGenerator[Union[ProgressMessage, StreamedTokens, AskModeResponse], None],
@@ -540,7 +546,10 @@ def run(
540546
context: Optional previous response from the agent.
541547
"""
542548
request_body = self._prepare_request_body(
543-
query=query, collections=collections, context=context
549+
query=query,
550+
collections=collections,
551+
context=context,
552+
result_evaluation="llm",
544553
)
545554

546555
response = httpx.post(
@@ -559,8 +568,11 @@ def ask(
559568
self,
560569
query: Union[str, list[ChatMessage]],
561570
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
571+
result_evaluation: Literal["llm", "none"] = "none",
562572
) -> AskModeResponse:
563-
request_body = self._prepare_request_body(query=query, collections=collections)
573+
request_body = self._prepare_request_body(
574+
query=query, collections=collections, result_evaluation=result_evaluation
575+
)
564576

565577
response = httpx.post(
566578
self.query_url + "/ask",
@@ -633,6 +645,7 @@ def stream(
633645
context=context,
634646
include_progress=include_progress,
635647
include_final_state=include_final_state,
648+
result_evaluation="llm",
636649
)
637650
with httpx.Client() as client:
638651
with connect_sse(
@@ -665,6 +678,7 @@ def ask_stream(
665678
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
666679
include_progress: Literal[True] = True,
667680
include_final_state: Literal[True] = True,
681+
result_evaluation: Literal["llm", "none"] = "none",
668682
) -> Generator[
669683
Union[ProgressMessage, StreamedTokens, AskModeResponse], None, None
670684
]: ...
@@ -676,6 +690,7 @@ def ask_stream(
676690
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
677691
include_progress: Literal[True] = True,
678692
include_final_state: Literal[False] = False,
693+
result_evaluation: Literal["llm", "none"] = "none",
679694
) -> Generator[Union[ProgressMessage, StreamedTokens], None, None]: ...
680695

681696
@overload
@@ -685,6 +700,7 @@ def ask_stream(
685700
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
686701
include_progress: Literal[False] = False,
687702
include_final_state: Literal[True] = True,
703+
result_evaluation: Literal["llm", "none"] = "none",
688704
) -> Generator[Union[StreamedTokens, AskModeResponse], None, None]: ...
689705

690706
@overload
@@ -694,6 +710,7 @@ def ask_stream(
694710
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
695711
include_progress: Literal[False] = False,
696712
include_final_state: Literal[False] = False,
713+
result_evaluation: Literal["llm", "none"] = "none",
697714
) -> Generator[StreamedTokens, None, None]: ...
698715

699716
def ask_stream(
@@ -702,13 +719,15 @@ def ask_stream(
702719
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
703720
include_progress: bool = True,
704721
include_final_state: bool = True,
722+
result_evaluation: Literal["llm", "none"] = "none",
705723
):
706724
"""Run the Query Agent ask mode and stream the response."""
707725
request_body = self._prepare_request_body(
708726
query=query,
709727
collections=collections,
710728
include_progress=include_progress,
711729
include_final_state=include_final_state,
730+
result_evaluation=result_evaluation,
712731
)
713732
with httpx.Client() as client:
714733
with connect_sse(
@@ -957,7 +976,10 @@ async def run(
957976
context: Optional previous response from the agent.
958977
"""
959978
request_body = self._prepare_request_body(
960-
query=query, collections=collections, context=context
979+
query=query,
980+
collections=collections,
981+
context=context,
982+
result_evaluation="llm",
961983
)
962984

963985
async with httpx.AsyncClient() as client:
@@ -977,8 +999,11 @@ async def ask(
977999
self,
9781000
query: Union[str, list[ChatMessage]],
9791001
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
1002+
result_evaluation: Literal["llm", "none"] = "none",
9801003
) -> AskModeResponse:
981-
request_body = self._prepare_request_body(query=query, collections=collections)
1004+
request_body = self._prepare_request_body(
1005+
query=query, collections=collections, result_evaluation=result_evaluation
1006+
)
9821007

9831008
async with httpx.AsyncClient() as client:
9841009
response = await client.post(
@@ -1053,6 +1078,7 @@ async def stream(
10531078
context=context,
10541079
include_progress=include_progress,
10551080
include_final_state=include_final_state,
1081+
result_evaluation="llm",
10561082
)
10571083
async with httpx.AsyncClient() as client:
10581084
async with aconnect_sse(
@@ -1085,6 +1111,7 @@ def ask_stream(
10851111
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
10861112
include_progress: Literal[True] = True,
10871113
include_final_state: Literal[True] = True,
1114+
result_evaluation: Literal["llm", "none"] = "none",
10881115
) -> AsyncGenerator[
10891116
Union[ProgressMessage, StreamedTokens, AskModeResponse], None
10901117
]: ...
@@ -1096,6 +1123,7 @@ def ask_stream(
10961123
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
10971124
include_progress: Literal[True] = True,
10981125
include_final_state: Literal[False] = False,
1126+
result_evaluation: Literal["llm", "none"] = "none",
10991127
) -> AsyncGenerator[Union[ProgressMessage, StreamedTokens], None]: ...
11001128

11011129
@overload
@@ -1105,6 +1133,7 @@ def ask_stream(
11051133
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
11061134
include_progress: Literal[False] = False,
11071135
include_final_state: Literal[True] = True,
1136+
result_evaluation: Literal["llm", "none"] = "none",
11081137
) -> AsyncGenerator[Union[StreamedTokens, AskModeResponse], None]: ...
11091138

11101139
@overload
@@ -1114,6 +1143,7 @@ def ask_stream(
11141143
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
11151144
include_progress: Literal[False] = False,
11161145
include_final_state: Literal[False] = False,
1146+
result_evaluation: Literal["llm", "none"] = "none",
11171147
) -> AsyncGenerator[StreamedTokens, None]: ...
11181148

11191149
async def ask_stream(
@@ -1122,13 +1152,15 @@ async def ask_stream(
11221152
collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
11231153
include_progress: bool = True,
11241154
include_final_state: bool = True,
1155+
result_evaluation: Literal["llm", "none"] = "none",
11251156
):
11261157
"""Run the Query Agent ask mode and stream the response."""
11271158
request_body = self._prepare_request_body(
11281159
query=query,
11291160
collections=collections,
11301161
include_progress=include_progress,
11311162
include_final_state=include_final_state,
1163+
result_evaluation=result_evaluation,
11321164
)
11331165
async with httpx.AsyncClient() as client:
11341166
async with aconnect_sse(

0 commit comments

Comments
 (0)