diff --git a/.release-please-manifest.json b/.release-please-manifest.json index cce9240..da59f99 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "0.3.1" + ".": "0.4.0" } \ No newline at end of file diff --git a/.stats.yml b/.stats.yml index f5f7eb3..e939ef2 100644 --- a/.stats.yml +++ b/.stats.yml @@ -1,4 +1,4 @@ configured_endpoints: 3 -openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/agentbase%2Fagentbase-b248631a2800e6d0f22a253717a63a551d9305e4c990dceb2f5c5361c519fa24.yml -openapi_spec_hash: 4c06f6155fb614add9b30cd294266e11 -config_hash: 2ad1177ae9ac1a57dc61c0f7e142fc06 +openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/agentbase%2Fagentbase-461274f650eccdf115b96ee1dbf1c0d4eb31485992cedc84ee0187aec72d63a5.yml +openapi_spec_hash: b8fa5fd0c89d9004d608af6061d0bf28 +config_hash: 8d9b688cf969a1760b1300b65b80e796 diff --git a/CHANGELOG.md b/CHANGELOG.md index f5e6cc6..9817ed5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## 0.4.0 (2025-10-11) + +Full Changelog: [v0.3.1...v0.4.0](https://github.com/AgentbaseHQ/agentbase-python/compare/v0.3.1...v0.4.0) + +### Features + +* **api:** manual updates ([0dd1219](https://github.com/AgentbaseHQ/agentbase-python/commit/0dd12194314a6f967ee3ca879e3663da2b75a911)) + ## 0.3.1 (2025-10-11) Full Changelog: [v0.3.0...v0.3.1](https://github.com/AgentbaseHQ/agentbase-python/compare/v0.3.0...v0.3.1) diff --git a/README.md b/README.md index d8a89a5..3722b75 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ It is generated with [Stainless](https://www.stainless.com/). ## Documentation -The full API of this library can be found in [api.md](api.md). +The REST API documentation can be found on [docs.agentbase.sh](https://docs.agentbase.sh). The full API of this library can be found in [api.md](api.md). ## Installation @@ -142,6 +142,22 @@ Nested request parameters are [TypedDicts](https://docs.python.org/3/library/typ Typed requests and responses provide autocomplete and documentation within your editor. If you would like to see type errors in VS Code to help catch bugs earlier, set `python.analysis.typeCheckingMode` to `basic`. +## Nested params + +Nested parameters are dictionaries, typed using `TypedDict`, for example: + +```python +from agentbase import Agentbase + +client = Agentbase() + +response = client.run_agent( + message="message", + callback={"url": "https://example.com"}, +) +print(response.callback) +``` + ## Handling errors When the library is unable to connect to the API (for example, due to network connection problems or a timeout), a subclass of `agentbase.APIConnectionError` is raised. diff --git a/SECURITY.md b/SECURITY.md index ec9ae63..1bd46b2 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -18,6 +18,10 @@ before making any information public. If you encounter security issues that are not directly related to SDKs but pertain to the services or products provided by Agentbase, please follow the respective company's security reporting guidelines. +### Agentbase Terms and Policies + +Please contact team@agentbase.sh for any questions or concerns regarding the security of our services. + --- Thank you for helping us keep the SDKs and systems they interact with secure. diff --git a/api.md b/api.md index c5c4ee0..4ddc008 100644 --- a/api.md +++ b/api.md @@ -10,26 +10,27 @@ Methods: - client.run_agent(\*\*params) -> str -# GetMessages +# Agent Types: ```python -from agentbase.types import GetMessageRetrieveResponse +from agentbase.types import AgentRunResponse ``` Methods: -- client.get_messages.retrieve(\*\*params) -> GetMessageRetrieveResponse +- client.agent.run(\*\*params) -> str -# ClearMessages +# Messages Types: ```python -from agentbase.types import ClearMessageClearResponse +from agentbase.types import MessageClearResponse, MessageGetResponse ``` Methods: -- client.clear_messages.clear(\*\*params) -> ClearMessageClearResponse +- client.messages.clear(\*\*params) -> MessageClearResponse +- client.messages.get(\*\*params) -> MessageGetResponse diff --git a/pyproject.toml b/pyproject.toml index 4181736..2f93d5b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,11 +1,11 @@ [project] name = "agentbase-sdk" -version = "0.3.1" +version = "0.4.0" description = "The official Python library for the agentbase API" dynamic = ["readme"] license = "Apache-2.0" authors = [ -{ name = "Agentbase", email = "" }, +{ name = "Agentbase", email = "team@agentbase.sh" }, ] dependencies = [ "httpx>=0.23.0, <1", diff --git a/src/agentbase/_client.py b/src/agentbase/_client.py index 4965dc2..d98a8e7 100644 --- a/src/agentbase/_client.py +++ b/src/agentbase/_client.py @@ -38,7 +38,7 @@ async_to_raw_response_wrapper, async_to_streamed_response_wrapper, ) -from .resources import get_messages, clear_messages +from .resources import agent, messages from ._streaming import Stream as Stream, AsyncStream as AsyncStream from ._exceptions import AgentbaseError, APIStatusError from ._base_client import ( @@ -62,8 +62,8 @@ class Agentbase(SyncAPIClient): - get_messages: get_messages.GetMessagesResource - clear_messages: clear_messages.ClearMessagesResource + agent: agent.AgentResource + messages: messages.MessagesResource with_raw_response: AgentbaseWithRawResponse with_streaming_response: AgentbaseWithStreamedResponse @@ -121,8 +121,8 @@ def __init__( _strict_response_validation=_strict_response_validation, ) - self.get_messages = get_messages.GetMessagesResource(self) - self.clear_messages = clear_messages.ClearMessagesResource(self) + self.agent = agent.AgentResource(self) + self.messages = messages.MessagesResource(self) self.with_raw_response = AgentbaseWithRawResponse(self) self.with_streaming_response = AgentbaseWithStreamedResponse(self) @@ -202,12 +202,17 @@ def run_agent( *, message: str, session: str | Omit = omit, + background: bool | Omit = omit, + callback: client_run_agent_params.Callback | Omit = omit, datastores: Iterable[client_run_agent_params.Datastore] | Omit = omit, + final_output: client_run_agent_params.FinalOutput | Omit = omit, mcp_servers: Iterable[client_run_agent_params.McpServer] | Omit = omit, mode: Literal["flash", "fast", "max"] | Omit = omit, + queries: Iterable[client_run_agent_params.Query] | Omit = omit, rules: SequenceNotStr[str] | Omit = omit, streaming_tokens: bool | Omit = omit, system: str | Omit = omit, + workflows: Iterable[client_run_agent_params.Workflow] | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: Headers | None = None, @@ -231,21 +236,38 @@ def run_agent( session: The session ID to continue the agent session conversation. If not provided, a new session will be created. + background: Whether to run the agent asynchronously on the server. When set to true, use + callback parameter to receive events. + + callback: A callback endpoint configuration to send agent message events back to. Use with + background true. + datastores: A set of datastores for the agent to utilize. Each object must include a `id` and `name`. + final_output: Configuration for an extra final output event that processes the entire agent + message thread and produces a structured output based on the provided JSON + schema. + mcp_servers: A list of MCP server configurations. Each object must include a `serverName` and `serverUrl`. mode: The agent mode. Allowed values are `flash`, `fast` or `max`. Defaults to `fast` if not supplied. + queries: A set of custom actions based on datastore (database) queries. Allows you to + quickly define actions that the agent can use to query your datastores. + rules: A list of constraints that the agent must follow. streaming_tokens: Whether to stream the agent messages token by token. system: A system prompt to provide system information to the agent. + workflows: A set of declarative workflows for the agent to execute. Each workflow is a DAG + (Directed Acyclic Graph) of steps that the agent interprets and executes + dynamically. + extra_headers: Send extra headers extra_query: Add additional query parameters to the request @@ -260,12 +282,17 @@ def run_agent( body=maybe_transform( { "message": message, + "background": background, + "callback": callback, "datastores": datastores, + "final_output": final_output, "mcp_servers": mcp_servers, "mode": mode, + "queries": queries, "rules": rules, "streaming_tokens": streaming_tokens, "system": system, + "workflows": workflows, }, client_run_agent_params.ClientRunAgentParams, ), @@ -316,8 +343,8 @@ def _make_status_error( class AsyncAgentbase(AsyncAPIClient): - get_messages: get_messages.AsyncGetMessagesResource - clear_messages: clear_messages.AsyncClearMessagesResource + agent: agent.AsyncAgentResource + messages: messages.AsyncMessagesResource with_raw_response: AsyncAgentbaseWithRawResponse with_streaming_response: AsyncAgentbaseWithStreamedResponse @@ -375,8 +402,8 @@ def __init__( _strict_response_validation=_strict_response_validation, ) - self.get_messages = get_messages.AsyncGetMessagesResource(self) - self.clear_messages = clear_messages.AsyncClearMessagesResource(self) + self.agent = agent.AsyncAgentResource(self) + self.messages = messages.AsyncMessagesResource(self) self.with_raw_response = AsyncAgentbaseWithRawResponse(self) self.with_streaming_response = AsyncAgentbaseWithStreamedResponse(self) @@ -456,12 +483,17 @@ async def run_agent( *, message: str, session: str | Omit = omit, + background: bool | Omit = omit, + callback: client_run_agent_params.Callback | Omit = omit, datastores: Iterable[client_run_agent_params.Datastore] | Omit = omit, + final_output: client_run_agent_params.FinalOutput | Omit = omit, mcp_servers: Iterable[client_run_agent_params.McpServer] | Omit = omit, mode: Literal["flash", "fast", "max"] | Omit = omit, + queries: Iterable[client_run_agent_params.Query] | Omit = omit, rules: SequenceNotStr[str] | Omit = omit, streaming_tokens: bool | Omit = omit, system: str | Omit = omit, + workflows: Iterable[client_run_agent_params.Workflow] | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: Headers | None = None, @@ -485,21 +517,38 @@ async def run_agent( session: The session ID to continue the agent session conversation. If not provided, a new session will be created. + background: Whether to run the agent asynchronously on the server. When set to true, use + callback parameter to receive events. + + callback: A callback endpoint configuration to send agent message events back to. Use with + background true. + datastores: A set of datastores for the agent to utilize. Each object must include a `id` and `name`. + final_output: Configuration for an extra final output event that processes the entire agent + message thread and produces a structured output based on the provided JSON + schema. + mcp_servers: A list of MCP server configurations. Each object must include a `serverName` and `serverUrl`. mode: The agent mode. Allowed values are `flash`, `fast` or `max`. Defaults to `fast` if not supplied. + queries: A set of custom actions based on datastore (database) queries. Allows you to + quickly define actions that the agent can use to query your datastores. + rules: A list of constraints that the agent must follow. streaming_tokens: Whether to stream the agent messages token by token. system: A system prompt to provide system information to the agent. + workflows: A set of declarative workflows for the agent to execute. Each workflow is a DAG + (Directed Acyclic Graph) of steps that the agent interprets and executes + dynamically. + extra_headers: Send extra headers extra_query: Add additional query parameters to the request @@ -514,12 +563,17 @@ async def run_agent( body=await async_maybe_transform( { "message": message, + "background": background, + "callback": callback, "datastores": datastores, + "final_output": final_output, "mcp_servers": mcp_servers, "mode": mode, + "queries": queries, "rules": rules, "streaming_tokens": streaming_tokens, "system": system, + "workflows": workflows, }, client_run_agent_params.ClientRunAgentParams, ), @@ -571,8 +625,8 @@ def _make_status_error( class AgentbaseWithRawResponse: def __init__(self, client: Agentbase) -> None: - self.get_messages = get_messages.GetMessagesResourceWithRawResponse(client.get_messages) - self.clear_messages = clear_messages.ClearMessagesResourceWithRawResponse(client.clear_messages) + self.agent = agent.AgentResourceWithRawResponse(client.agent) + self.messages = messages.MessagesResourceWithRawResponse(client.messages) self.run_agent = to_raw_response_wrapper( client.run_agent, @@ -581,8 +635,8 @@ def __init__(self, client: Agentbase) -> None: class AsyncAgentbaseWithRawResponse: def __init__(self, client: AsyncAgentbase) -> None: - self.get_messages = get_messages.AsyncGetMessagesResourceWithRawResponse(client.get_messages) - self.clear_messages = clear_messages.AsyncClearMessagesResourceWithRawResponse(client.clear_messages) + self.agent = agent.AsyncAgentResourceWithRawResponse(client.agent) + self.messages = messages.AsyncMessagesResourceWithRawResponse(client.messages) self.run_agent = async_to_raw_response_wrapper( client.run_agent, @@ -591,8 +645,8 @@ def __init__(self, client: AsyncAgentbase) -> None: class AgentbaseWithStreamedResponse: def __init__(self, client: Agentbase) -> None: - self.get_messages = get_messages.GetMessagesResourceWithStreamingResponse(client.get_messages) - self.clear_messages = clear_messages.ClearMessagesResourceWithStreamingResponse(client.clear_messages) + self.agent = agent.AgentResourceWithStreamingResponse(client.agent) + self.messages = messages.MessagesResourceWithStreamingResponse(client.messages) self.run_agent = to_streamed_response_wrapper( client.run_agent, @@ -601,8 +655,8 @@ def __init__(self, client: Agentbase) -> None: class AsyncAgentbaseWithStreamedResponse: def __init__(self, client: AsyncAgentbase) -> None: - self.get_messages = get_messages.AsyncGetMessagesResourceWithStreamingResponse(client.get_messages) - self.clear_messages = clear_messages.AsyncClearMessagesResourceWithStreamingResponse(client.clear_messages) + self.agent = agent.AsyncAgentResourceWithStreamingResponse(client.agent) + self.messages = messages.AsyncMessagesResourceWithStreamingResponse(client.messages) self.run_agent = async_to_streamed_response_wrapper( client.run_agent, diff --git a/src/agentbase/_version.py b/src/agentbase/_version.py index fdd20c1..87bf175 100644 --- a/src/agentbase/_version.py +++ b/src/agentbase/_version.py @@ -1,4 +1,4 @@ # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. __title__ = "agentbase" -__version__ = "0.3.1" # x-release-please-version +__version__ = "0.4.0" # x-release-please-version diff --git a/src/agentbase/resources/__init__.py b/src/agentbase/resources/__init__.py index 6348735..5da984c 100644 --- a/src/agentbase/resources/__init__.py +++ b/src/agentbase/resources/__init__.py @@ -1,33 +1,33 @@ # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. -from .get_messages import ( - GetMessagesResource, - AsyncGetMessagesResource, - GetMessagesResourceWithRawResponse, - AsyncGetMessagesResourceWithRawResponse, - GetMessagesResourceWithStreamingResponse, - AsyncGetMessagesResourceWithStreamingResponse, +from .agent import ( + AgentResource, + AsyncAgentResource, + AgentResourceWithRawResponse, + AsyncAgentResourceWithRawResponse, + AgentResourceWithStreamingResponse, + AsyncAgentResourceWithStreamingResponse, ) -from .clear_messages import ( - ClearMessagesResource, - AsyncClearMessagesResource, - ClearMessagesResourceWithRawResponse, - AsyncClearMessagesResourceWithRawResponse, - ClearMessagesResourceWithStreamingResponse, - AsyncClearMessagesResourceWithStreamingResponse, +from .messages import ( + MessagesResource, + AsyncMessagesResource, + MessagesResourceWithRawResponse, + AsyncMessagesResourceWithRawResponse, + MessagesResourceWithStreamingResponse, + AsyncMessagesResourceWithStreamingResponse, ) __all__ = [ - "GetMessagesResource", - "AsyncGetMessagesResource", - "GetMessagesResourceWithRawResponse", - "AsyncGetMessagesResourceWithRawResponse", - "GetMessagesResourceWithStreamingResponse", - "AsyncGetMessagesResourceWithStreamingResponse", - "ClearMessagesResource", - "AsyncClearMessagesResource", - "ClearMessagesResourceWithRawResponse", - "AsyncClearMessagesResourceWithRawResponse", - "ClearMessagesResourceWithStreamingResponse", - "AsyncClearMessagesResourceWithStreamingResponse", + "AgentResource", + "AsyncAgentResource", + "AgentResourceWithRawResponse", + "AsyncAgentResourceWithRawResponse", + "AgentResourceWithStreamingResponse", + "AsyncAgentResourceWithStreamingResponse", + "MessagesResource", + "AsyncMessagesResource", + "MessagesResourceWithRawResponse", + "AsyncMessagesResourceWithRawResponse", + "MessagesResourceWithStreamingResponse", + "AsyncMessagesResourceWithStreamingResponse", ] diff --git a/src/agentbase/resources/agent.py b/src/agentbase/resources/agent.py new file mode 100644 index 0000000..51a3703 --- /dev/null +++ b/src/agentbase/resources/agent.py @@ -0,0 +1,327 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from __future__ import annotations + +from typing import Iterable +from typing_extensions import Literal + +import httpx + +from ..types import agent_run_params +from .._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given +from .._utils import maybe_transform, async_maybe_transform +from .._compat import cached_property +from .._resource import SyncAPIResource, AsyncAPIResource +from .._response import ( + to_raw_response_wrapper, + to_streamed_response_wrapper, + async_to_raw_response_wrapper, + async_to_streamed_response_wrapper, +) +from .._streaming import Stream, AsyncStream +from .._base_client import make_request_options +from ..types.agent_run_response import AgentRunResponse + +__all__ = ["AgentResource", "AsyncAgentResource"] + + +class AgentResource(SyncAPIResource): + @cached_property + def with_raw_response(self) -> AgentResourceWithRawResponse: + """ + This property can be used as a prefix for any HTTP method call to return + the raw response object instead of the parsed content. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers + """ + return AgentResourceWithRawResponse(self) + + @cached_property + def with_streaming_response(self) -> AgentResourceWithStreamingResponse: + """ + An alternative to `.with_raw_response` that doesn't eagerly read the response body. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response + """ + return AgentResourceWithStreamingResponse(self) + + def run( + self, + *, + message: str, + session: str | Omit = omit, + background: bool | Omit = omit, + callback: agent_run_params.Callback | Omit = omit, + datastores: Iterable[agent_run_params.Datastore] | Omit = omit, + final_output: agent_run_params.FinalOutput | Omit = omit, + mcp_servers: Iterable[agent_run_params.McpServer] | Omit = omit, + mode: Literal["flash", "fast", "max"] | Omit = omit, + queries: Iterable[agent_run_params.Query] | Omit = omit, + rules: SequenceNotStr[str] | Omit = omit, + streaming_tokens: bool | Omit = omit, + system: str | Omit = omit, + workflows: Iterable[agent_run_params.Workflow] | Omit = omit, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> Stream[AgentRunResponse]: + """Run an agent on a task or message. + + + A new session can be created by omitting the `session` query parameter, or an + existing session can be continued by specifying the session ID in the `session` + query parameter. + The request body includes the message, optional system prompt, mode, MCP server + configuration, optional rules and whether the response should be streamed. + The response is a streaming response and returns a sequence of events + representing the agent’s thoughts and responses. + + Args: + message: The task or message to run the agent with. + + session: The session ID to continue the agent session conversation. If not provided, a + new session will be created. + + background: Whether to run the agent asynchronously on the server. When set to true, use + callback parameter to receive events. + + callback: A callback endpoint configuration to send agent message events back to. Use with + background true. + + datastores: A set of datastores for the agent to utilize. Each object must include a `id` + and `name`. + + final_output: Configuration for an extra final output event that processes the entire agent + message thread and produces a structured output based on the provided JSON + schema. + + mcp_servers: A list of MCP server configurations. Each object must include a `serverName` and + `serverUrl`. + + mode: The agent mode. Allowed values are `flash`, `fast` or `max`. Defaults to `fast` + if not supplied. + + queries: A set of custom actions based on datastore (database) queries. Allows you to + quickly define actions that the agent can use to query your datastores. + + rules: A list of constraints that the agent must follow. + + streaming_tokens: Whether to stream the agent messages token by token. + + system: A system prompt to provide system information to the agent. + + workflows: A set of declarative workflows for the agent to execute. Each workflow is a DAG + (Directed Acyclic Graph) of steps that the agent interprets and executes + dynamically. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + timeout: Override the client-level default timeout for this request, in seconds + """ + extra_headers = {"Accept": "text/event-stream", **(extra_headers or {})} + return self._post( + "/", + body=maybe_transform( + { + "message": message, + "background": background, + "callback": callback, + "datastores": datastores, + "final_output": final_output, + "mcp_servers": mcp_servers, + "mode": mode, + "queries": queries, + "rules": rules, + "streaming_tokens": streaming_tokens, + "system": system, + "workflows": workflows, + }, + agent_run_params.AgentRunParams, + ), + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + query=maybe_transform({"session": session}, agent_run_params.AgentRunParams), + ), + cast_to=str, + stream=True, + stream_cls=Stream[AgentRunResponse], + ) + + +class AsyncAgentResource(AsyncAPIResource): + @cached_property + def with_raw_response(self) -> AsyncAgentResourceWithRawResponse: + """ + This property can be used as a prefix for any HTTP method call to return + the raw response object instead of the parsed content. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers + """ + return AsyncAgentResourceWithRawResponse(self) + + @cached_property + def with_streaming_response(self) -> AsyncAgentResourceWithStreamingResponse: + """ + An alternative to `.with_raw_response` that doesn't eagerly read the response body. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response + """ + return AsyncAgentResourceWithStreamingResponse(self) + + async def run( + self, + *, + message: str, + session: str | Omit = omit, + background: bool | Omit = omit, + callback: agent_run_params.Callback | Omit = omit, + datastores: Iterable[agent_run_params.Datastore] | Omit = omit, + final_output: agent_run_params.FinalOutput | Omit = omit, + mcp_servers: Iterable[agent_run_params.McpServer] | Omit = omit, + mode: Literal["flash", "fast", "max"] | Omit = omit, + queries: Iterable[agent_run_params.Query] | Omit = omit, + rules: SequenceNotStr[str] | Omit = omit, + streaming_tokens: bool | Omit = omit, + system: str | Omit = omit, + workflows: Iterable[agent_run_params.Workflow] | Omit = omit, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> AsyncStream[AgentRunResponse]: + """Run an agent on a task or message. + + + A new session can be created by omitting the `session` query parameter, or an + existing session can be continued by specifying the session ID in the `session` + query parameter. + The request body includes the message, optional system prompt, mode, MCP server + configuration, optional rules and whether the response should be streamed. + The response is a streaming response and returns a sequence of events + representing the agent’s thoughts and responses. + + Args: + message: The task or message to run the agent with. + + session: The session ID to continue the agent session conversation. If not provided, a + new session will be created. + + background: Whether to run the agent asynchronously on the server. When set to true, use + callback parameter to receive events. + + callback: A callback endpoint configuration to send agent message events back to. Use with + background true. + + datastores: A set of datastores for the agent to utilize. Each object must include a `id` + and `name`. + + final_output: Configuration for an extra final output event that processes the entire agent + message thread and produces a structured output based on the provided JSON + schema. + + mcp_servers: A list of MCP server configurations. Each object must include a `serverName` and + `serverUrl`. + + mode: The agent mode. Allowed values are `flash`, `fast` or `max`. Defaults to `fast` + if not supplied. + + queries: A set of custom actions based on datastore (database) queries. Allows you to + quickly define actions that the agent can use to query your datastores. + + rules: A list of constraints that the agent must follow. + + streaming_tokens: Whether to stream the agent messages token by token. + + system: A system prompt to provide system information to the agent. + + workflows: A set of declarative workflows for the agent to execute. Each workflow is a DAG + (Directed Acyclic Graph) of steps that the agent interprets and executes + dynamically. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + timeout: Override the client-level default timeout for this request, in seconds + """ + extra_headers = {"Accept": "text/event-stream", **(extra_headers or {})} + return await self._post( + "/", + body=await async_maybe_transform( + { + "message": message, + "background": background, + "callback": callback, + "datastores": datastores, + "final_output": final_output, + "mcp_servers": mcp_servers, + "mode": mode, + "queries": queries, + "rules": rules, + "streaming_tokens": streaming_tokens, + "system": system, + "workflows": workflows, + }, + agent_run_params.AgentRunParams, + ), + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + query=await async_maybe_transform({"session": session}, agent_run_params.AgentRunParams), + ), + cast_to=str, + stream=True, + stream_cls=AsyncStream[AgentRunResponse], + ) + + +class AgentResourceWithRawResponse: + def __init__(self, agent: AgentResource) -> None: + self._agent = agent + + self.run = to_raw_response_wrapper( + agent.run, + ) + + +class AsyncAgentResourceWithRawResponse: + def __init__(self, agent: AsyncAgentResource) -> None: + self._agent = agent + + self.run = async_to_raw_response_wrapper( + agent.run, + ) + + +class AgentResourceWithStreamingResponse: + def __init__(self, agent: AgentResource) -> None: + self._agent = agent + + self.run = to_streamed_response_wrapper( + agent.run, + ) + + +class AsyncAgentResourceWithStreamingResponse: + def __init__(self, agent: AsyncAgentResource) -> None: + self._agent = agent + + self.run = async_to_streamed_response_wrapper( + agent.run, + ) diff --git a/src/agentbase/resources/clear_messages.py b/src/agentbase/resources/clear_messages.py deleted file mode 100644 index 2bc95d9..0000000 --- a/src/agentbase/resources/clear_messages.py +++ /dev/null @@ -1,175 +0,0 @@ -# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. - -from __future__ import annotations - -import httpx - -from ..types import clear_message_clear_params -from .._types import Body, Query, Headers, NotGiven, not_given -from .._utils import maybe_transform, async_maybe_transform -from .._compat import cached_property -from .._resource import SyncAPIResource, AsyncAPIResource -from .._response import ( - to_raw_response_wrapper, - to_streamed_response_wrapper, - async_to_raw_response_wrapper, - async_to_streamed_response_wrapper, -) -from .._base_client import make_request_options -from ..types.clear_message_clear_response import ClearMessageClearResponse - -__all__ = ["ClearMessagesResource", "AsyncClearMessagesResource"] - - -class ClearMessagesResource(SyncAPIResource): - @cached_property - def with_raw_response(self) -> ClearMessagesResourceWithRawResponse: - """ - This property can be used as a prefix for any HTTP method call to return - the raw response object instead of the parsed content. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers - """ - return ClearMessagesResourceWithRawResponse(self) - - @cached_property - def with_streaming_response(self) -> ClearMessagesResourceWithStreamingResponse: - """ - An alternative to `.with_raw_response` that doesn't eagerly read the response body. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response - """ - return ClearMessagesResourceWithStreamingResponse(self) - - def clear( - self, - *, - session: str, - # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. - # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Headers | None = None, - extra_query: Query | None = None, - extra_body: Body | None = None, - timeout: float | httpx.Timeout | None | NotGiven = not_given, - ) -> ClearMessageClearResponse: - """ - Clear all messages from a given agent session. - - Args: - session: The session ID to clear messages from. - - extra_headers: Send extra headers - - extra_query: Add additional query parameters to the request - - extra_body: Add additional JSON properties to the request - - timeout: Override the client-level default timeout for this request, in seconds - """ - return self._post( - "/clear-messages", - options=make_request_options( - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - query=maybe_transform({"session": session}, clear_message_clear_params.ClearMessageClearParams), - ), - cast_to=ClearMessageClearResponse, - ) - - -class AsyncClearMessagesResource(AsyncAPIResource): - @cached_property - def with_raw_response(self) -> AsyncClearMessagesResourceWithRawResponse: - """ - This property can be used as a prefix for any HTTP method call to return - the raw response object instead of the parsed content. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers - """ - return AsyncClearMessagesResourceWithRawResponse(self) - - @cached_property - def with_streaming_response(self) -> AsyncClearMessagesResourceWithStreamingResponse: - """ - An alternative to `.with_raw_response` that doesn't eagerly read the response body. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response - """ - return AsyncClearMessagesResourceWithStreamingResponse(self) - - async def clear( - self, - *, - session: str, - # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. - # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Headers | None = None, - extra_query: Query | None = None, - extra_body: Body | None = None, - timeout: float | httpx.Timeout | None | NotGiven = not_given, - ) -> ClearMessageClearResponse: - """ - Clear all messages from a given agent session. - - Args: - session: The session ID to clear messages from. - - extra_headers: Send extra headers - - extra_query: Add additional query parameters to the request - - extra_body: Add additional JSON properties to the request - - timeout: Override the client-level default timeout for this request, in seconds - """ - return await self._post( - "/clear-messages", - options=make_request_options( - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - query=await async_maybe_transform( - {"session": session}, clear_message_clear_params.ClearMessageClearParams - ), - ), - cast_to=ClearMessageClearResponse, - ) - - -class ClearMessagesResourceWithRawResponse: - def __init__(self, clear_messages: ClearMessagesResource) -> None: - self._clear_messages = clear_messages - - self.clear = to_raw_response_wrapper( - clear_messages.clear, - ) - - -class AsyncClearMessagesResourceWithRawResponse: - def __init__(self, clear_messages: AsyncClearMessagesResource) -> None: - self._clear_messages = clear_messages - - self.clear = async_to_raw_response_wrapper( - clear_messages.clear, - ) - - -class ClearMessagesResourceWithStreamingResponse: - def __init__(self, clear_messages: ClearMessagesResource) -> None: - self._clear_messages = clear_messages - - self.clear = to_streamed_response_wrapper( - clear_messages.clear, - ) - - -class AsyncClearMessagesResourceWithStreamingResponse: - def __init__(self, clear_messages: AsyncClearMessagesResource) -> None: - self._clear_messages = clear_messages - - self.clear = async_to_streamed_response_wrapper( - clear_messages.clear, - ) diff --git a/src/agentbase/resources/get_messages.py b/src/agentbase/resources/get_messages.py deleted file mode 100644 index d47a0db..0000000 --- a/src/agentbase/resources/get_messages.py +++ /dev/null @@ -1,181 +0,0 @@ -# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. - -from __future__ import annotations - -import httpx - -from ..types import get_message_retrieve_params -from .._types import Body, Query, Headers, NotGiven, not_given -from .._utils import maybe_transform, async_maybe_transform -from .._compat import cached_property -from .._resource import SyncAPIResource, AsyncAPIResource -from .._response import ( - to_raw_response_wrapper, - to_streamed_response_wrapper, - async_to_raw_response_wrapper, - async_to_streamed_response_wrapper, -) -from .._base_client import make_request_options -from ..types.get_message_retrieve_response import GetMessageRetrieveResponse - -__all__ = ["GetMessagesResource", "AsyncGetMessagesResource"] - - -class GetMessagesResource(SyncAPIResource): - @cached_property - def with_raw_response(self) -> GetMessagesResourceWithRawResponse: - """ - This property can be used as a prefix for any HTTP method call to return - the raw response object instead of the parsed content. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers - """ - return GetMessagesResourceWithRawResponse(self) - - @cached_property - def with_streaming_response(self) -> GetMessagesResourceWithStreamingResponse: - """ - An alternative to `.with_raw_response` that doesn't eagerly read the response body. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response - """ - return GetMessagesResourceWithStreamingResponse(self) - - def retrieve( - self, - *, - session: str, - # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. - # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Headers | None = None, - extra_query: Query | None = None, - extra_body: Body | None = None, - timeout: float | httpx.Timeout | None | NotGiven = not_given, - ) -> GetMessageRetrieveResponse: - """Retrieve the entire message history for a given agent session. - - - Messages include user messages, the agent’s internal thoughts, agent responses, - and tool usage records. - - Args: - session: The session ID to retrieve messages from. - - extra_headers: Send extra headers - - extra_query: Add additional query parameters to the request - - extra_body: Add additional JSON properties to the request - - timeout: Override the client-level default timeout for this request, in seconds - """ - return self._post( - "/get-messages", - options=make_request_options( - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - query=maybe_transform({"session": session}, get_message_retrieve_params.GetMessageRetrieveParams), - ), - cast_to=GetMessageRetrieveResponse, - ) - - -class AsyncGetMessagesResource(AsyncAPIResource): - @cached_property - def with_raw_response(self) -> AsyncGetMessagesResourceWithRawResponse: - """ - This property can be used as a prefix for any HTTP method call to return - the raw response object instead of the parsed content. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers - """ - return AsyncGetMessagesResourceWithRawResponse(self) - - @cached_property - def with_streaming_response(self) -> AsyncGetMessagesResourceWithStreamingResponse: - """ - An alternative to `.with_raw_response` that doesn't eagerly read the response body. - - For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response - """ - return AsyncGetMessagesResourceWithStreamingResponse(self) - - async def retrieve( - self, - *, - session: str, - # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. - # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Headers | None = None, - extra_query: Query | None = None, - extra_body: Body | None = None, - timeout: float | httpx.Timeout | None | NotGiven = not_given, - ) -> GetMessageRetrieveResponse: - """Retrieve the entire message history for a given agent session. - - - Messages include user messages, the agent’s internal thoughts, agent responses, - and tool usage records. - - Args: - session: The session ID to retrieve messages from. - - extra_headers: Send extra headers - - extra_query: Add additional query parameters to the request - - extra_body: Add additional JSON properties to the request - - timeout: Override the client-level default timeout for this request, in seconds - """ - return await self._post( - "/get-messages", - options=make_request_options( - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - query=await async_maybe_transform( - {"session": session}, get_message_retrieve_params.GetMessageRetrieveParams - ), - ), - cast_to=GetMessageRetrieveResponse, - ) - - -class GetMessagesResourceWithRawResponse: - def __init__(self, get_messages: GetMessagesResource) -> None: - self._get_messages = get_messages - - self.retrieve = to_raw_response_wrapper( - get_messages.retrieve, - ) - - -class AsyncGetMessagesResourceWithRawResponse: - def __init__(self, get_messages: AsyncGetMessagesResource) -> None: - self._get_messages = get_messages - - self.retrieve = async_to_raw_response_wrapper( - get_messages.retrieve, - ) - - -class GetMessagesResourceWithStreamingResponse: - def __init__(self, get_messages: GetMessagesResource) -> None: - self._get_messages = get_messages - - self.retrieve = to_streamed_response_wrapper( - get_messages.retrieve, - ) - - -class AsyncGetMessagesResourceWithStreamingResponse: - def __init__(self, get_messages: AsyncGetMessagesResource) -> None: - self._get_messages = get_messages - - self.retrieve = async_to_streamed_response_wrapper( - get_messages.retrieve, - ) diff --git a/src/agentbase/resources/messages.py b/src/agentbase/resources/messages.py new file mode 100644 index 0000000..345a886 --- /dev/null +++ b/src/agentbase/resources/messages.py @@ -0,0 +1,266 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from __future__ import annotations + +import httpx + +from ..types import message_get_params, message_clear_params +from .._types import Body, Query, Headers, NotGiven, not_given +from .._utils import maybe_transform, async_maybe_transform +from .._compat import cached_property +from .._resource import SyncAPIResource, AsyncAPIResource +from .._response import ( + to_raw_response_wrapper, + to_streamed_response_wrapper, + async_to_raw_response_wrapper, + async_to_streamed_response_wrapper, +) +from .._base_client import make_request_options +from ..types.message_get_response import MessageGetResponse +from ..types.message_clear_response import MessageClearResponse + +__all__ = ["MessagesResource", "AsyncMessagesResource"] + + +class MessagesResource(SyncAPIResource): + @cached_property + def with_raw_response(self) -> MessagesResourceWithRawResponse: + """ + This property can be used as a prefix for any HTTP method call to return + the raw response object instead of the parsed content. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers + """ + return MessagesResourceWithRawResponse(self) + + @cached_property + def with_streaming_response(self) -> MessagesResourceWithStreamingResponse: + """ + An alternative to `.with_raw_response` that doesn't eagerly read the response body. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response + """ + return MessagesResourceWithStreamingResponse(self) + + def clear( + self, + *, + session: str, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> MessageClearResponse: + """ + Clear all messages from a given agent session. + + Args: + session: The session ID to clear messages from. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + timeout: Override the client-level default timeout for this request, in seconds + """ + return self._post( + "/clear-messages", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + query=maybe_transform({"session": session}, message_clear_params.MessageClearParams), + ), + cast_to=MessageClearResponse, + ) + + def get( + self, + *, + session: str, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> MessageGetResponse: + """Retrieve the entire message history for a given agent session. + + + Messages include user messages, the agent’s internal thoughts, agent responses, + and tool usage records. + + Args: + session: The session ID to retrieve messages from. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + timeout: Override the client-level default timeout for this request, in seconds + """ + return self._post( + "/get-messages", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + query=maybe_transform({"session": session}, message_get_params.MessageGetParams), + ), + cast_to=MessageGetResponse, + ) + + +class AsyncMessagesResource(AsyncAPIResource): + @cached_property + def with_raw_response(self) -> AsyncMessagesResourceWithRawResponse: + """ + This property can be used as a prefix for any HTTP method call to return + the raw response object instead of the parsed content. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#accessing-raw-response-data-eg-headers + """ + return AsyncMessagesResourceWithRawResponse(self) + + @cached_property + def with_streaming_response(self) -> AsyncMessagesResourceWithStreamingResponse: + """ + An alternative to `.with_raw_response` that doesn't eagerly read the response body. + + For more information, see https://www.github.com/AgentbaseHQ/agentbase-python#with_streaming_response + """ + return AsyncMessagesResourceWithStreamingResponse(self) + + async def clear( + self, + *, + session: str, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> MessageClearResponse: + """ + Clear all messages from a given agent session. + + Args: + session: The session ID to clear messages from. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + timeout: Override the client-level default timeout for this request, in seconds + """ + return await self._post( + "/clear-messages", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + query=await async_maybe_transform({"session": session}, message_clear_params.MessageClearParams), + ), + cast_to=MessageClearResponse, + ) + + async def get( + self, + *, + session: str, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> MessageGetResponse: + """Retrieve the entire message history for a given agent session. + + + Messages include user messages, the agent’s internal thoughts, agent responses, + and tool usage records. + + Args: + session: The session ID to retrieve messages from. + + extra_headers: Send extra headers + + extra_query: Add additional query parameters to the request + + extra_body: Add additional JSON properties to the request + + timeout: Override the client-level default timeout for this request, in seconds + """ + return await self._post( + "/get-messages", + options=make_request_options( + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + query=await async_maybe_transform({"session": session}, message_get_params.MessageGetParams), + ), + cast_to=MessageGetResponse, + ) + + +class MessagesResourceWithRawResponse: + def __init__(self, messages: MessagesResource) -> None: + self._messages = messages + + self.clear = to_raw_response_wrapper( + messages.clear, + ) + self.get = to_raw_response_wrapper( + messages.get, + ) + + +class AsyncMessagesResourceWithRawResponse: + def __init__(self, messages: AsyncMessagesResource) -> None: + self._messages = messages + + self.clear = async_to_raw_response_wrapper( + messages.clear, + ) + self.get = async_to_raw_response_wrapper( + messages.get, + ) + + +class MessagesResourceWithStreamingResponse: + def __init__(self, messages: MessagesResource) -> None: + self._messages = messages + + self.clear = to_streamed_response_wrapper( + messages.clear, + ) + self.get = to_streamed_response_wrapper( + messages.get, + ) + + +class AsyncMessagesResourceWithStreamingResponse: + def __init__(self, messages: AsyncMessagesResource) -> None: + self._messages = messages + + self.clear = async_to_streamed_response_wrapper( + messages.clear, + ) + self.get = async_to_streamed_response_wrapper( + messages.get, + ) diff --git a/src/agentbase/types/__init__.py b/src/agentbase/types/__init__.py index 7d13cb4..ca25e6e 100644 --- a/src/agentbase/types/__init__.py +++ b/src/agentbase/types/__init__.py @@ -2,9 +2,11 @@ from __future__ import annotations +from .agent_run_params import AgentRunParams as AgentRunParams +from .agent_run_response import AgentRunResponse as AgentRunResponse +from .message_get_params import MessageGetParams as MessageGetParams from .run_agent_response import RunAgentResponse as RunAgentResponse +from .message_clear_params import MessageClearParams as MessageClearParams +from .message_get_response import MessageGetResponse as MessageGetResponse +from .message_clear_response import MessageClearResponse as MessageClearResponse from .client_run_agent_params import ClientRunAgentParams as ClientRunAgentParams -from .clear_message_clear_params import ClearMessageClearParams as ClearMessageClearParams -from .get_message_retrieve_params import GetMessageRetrieveParams as GetMessageRetrieveParams -from .clear_message_clear_response import ClearMessageClearResponse as ClearMessageClearResponse -from .get_message_retrieve_response import GetMessageRetrieveResponse as GetMessageRetrieveResponse diff --git a/src/agentbase/types/agent_run_params.py b/src/agentbase/types/agent_run_params.py new file mode 100644 index 0000000..63fbbdd --- /dev/null +++ b/src/agentbase/types/agent_run_params.py @@ -0,0 +1,181 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from __future__ import annotations + +from typing import Dict, Iterable +from typing_extensions import Literal, Required, Annotated, TypedDict + +from .._types import SequenceNotStr +from .._utils import PropertyInfo + +__all__ = [ + "AgentRunParams", + "Callback", + "Datastore", + "FinalOutput", + "McpServer", + "Query", + "Workflow", + "WorkflowStep", + "WorkflowStepRetryPolicy", +] + + +class AgentRunParams(TypedDict, total=False): + message: Required[str] + """The task or message to run the agent with.""" + + session: str + """The session ID to continue the agent session conversation. + + If not provided, a new session will be created. + """ + + background: bool + """Whether to run the agent asynchronously on the server. + + When set to true, use callback parameter to receive events. + """ + + callback: Callback + """A callback endpoint configuration to send agent message events back to. + + Use with background true. + """ + + datastores: Iterable[Datastore] + """A set of datastores for the agent to utilize. + + Each object must include a `id` and `name`. + """ + + final_output: FinalOutput + """ + Configuration for an extra final output event that processes the entire agent + message thread and produces a structured output based on the provided JSON + schema. + """ + + mcp_servers: Iterable[McpServer] + """A list of MCP server configurations. + + Each object must include a `serverName` and `serverUrl`. + """ + + mode: Literal["flash", "fast", "max"] + """The agent mode. + + Allowed values are `flash`, `fast` or `max`. Defaults to `fast` if not supplied. + """ + + queries: Iterable[Query] + """A set of custom actions based on datastore (database) queries. + + Allows you to quickly define actions that the agent can use to query your + datastores. + """ + + rules: SequenceNotStr[str] + """A list of constraints that the agent must follow.""" + + streaming_tokens: bool + """Whether to stream the agent messages token by token.""" + + system: str + """A system prompt to provide system information to the agent.""" + + workflows: Iterable[Workflow] + """A set of declarative workflows for the agent to execute. + + Each workflow is a DAG (Directed Acyclic Graph) of steps that the agent + interprets and executes dynamically. + """ + + +class Callback(TypedDict, total=False): + url: Required[str] + """The webhook URL to send events to.""" + + headers: Dict[str, str] + """Optional headers to include in the callback request.""" + + +class Datastore(TypedDict, total=False): + id: Required[str] + """The ID of the datastore.""" + + name: Required[str] + """The name of the datastore.""" + + +class FinalOutput(TypedDict, total=False): + name: Required[str] + """Name for the final output.""" + + schema: Required[object] + """JSON schema defining the structure of the final output.""" + + strict: bool + """Whether to enforce strict schema validation.""" + + +class McpServer(TypedDict, total=False): + server_name: Required[Annotated[str, PropertyInfo(alias="serverName")]] + """Name of the MCP server.""" + + server_url: Required[Annotated[str, PropertyInfo(alias="serverUrl")]] + """URL of the MCP server.""" + + +class Query(TypedDict, total=False): + description: Required[str] + """Description of what the query does.""" + + name: Required[str] + """Name of the query action.""" + + query: Required[str] + """The SQL query to execute.""" + + +class WorkflowStepRetryPolicy(TypedDict, total=False): + backoff: str + + max_attempts: int + + +class WorkflowStep(TypedDict, total=False): + id: Required[str] + """Unique identifier for the step.""" + + depends_on: Required[SequenceNotStr[str]] + """Array of step IDs that must complete before this step runs.""" + + description: Required[str] + """What the step should accomplish.""" + + name: Required[str] + """Name of the step.""" + + optional: bool + """Whether the step can be skipped if it fails.""" + + output_schema: object + """JSON schema for expected output validation.""" + + retry_policy: WorkflowStepRetryPolicy + """Retry configuration for the step.""" + + +class Workflow(TypedDict, total=False): + id: Required[str] + """Unique identifier for the workflow.""" + + description: Required[str] + """What the workflow accomplishes.""" + + name: Required[str] + """Name of the workflow.""" + + steps: Required[Iterable[WorkflowStep]] + """Array of step objects.""" diff --git a/src/agentbase/types/agent_run_response.py b/src/agentbase/types/agent_run_response.py new file mode 100644 index 0000000..7d7a853 --- /dev/null +++ b/src/agentbase/types/agent_run_response.py @@ -0,0 +1,7 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing_extensions import TypeAlias + +__all__ = ["AgentRunResponse"] + +AgentRunResponse: TypeAlias = str diff --git a/src/agentbase/types/client_run_agent_params.py b/src/agentbase/types/client_run_agent_params.py index 75a4d84..1143bc8 100644 --- a/src/agentbase/types/client_run_agent_params.py +++ b/src/agentbase/types/client_run_agent_params.py @@ -2,13 +2,23 @@ from __future__ import annotations -from typing import Iterable +from typing import Dict, Iterable from typing_extensions import Literal, Required, Annotated, TypedDict from .._types import SequenceNotStr from .._utils import PropertyInfo -__all__ = ["ClientRunAgentParams", "Datastore", "McpServer"] +__all__ = [ + "ClientRunAgentParams", + "Callback", + "Datastore", + "FinalOutput", + "McpServer", + "Query", + "Workflow", + "WorkflowStep", + "WorkflowStepRetryPolicy", +] class ClientRunAgentParams(TypedDict, total=False): @@ -21,12 +31,31 @@ class ClientRunAgentParams(TypedDict, total=False): If not provided, a new session will be created. """ + background: bool + """Whether to run the agent asynchronously on the server. + + When set to true, use callback parameter to receive events. + """ + + callback: Callback + """A callback endpoint configuration to send agent message events back to. + + Use with background true. + """ + datastores: Iterable[Datastore] """A set of datastores for the agent to utilize. Each object must include a `id` and `name`. """ + final_output: FinalOutput + """ + Configuration for an extra final output event that processes the entire agent + message thread and produces a structured output based on the provided JSON + schema. + """ + mcp_servers: Iterable[McpServer] """A list of MCP server configurations. @@ -39,6 +68,13 @@ class ClientRunAgentParams(TypedDict, total=False): Allowed values are `flash`, `fast` or `max`. Defaults to `fast` if not supplied. """ + queries: Iterable[Query] + """A set of custom actions based on datastore (database) queries. + + Allows you to quickly define actions that the agent can use to query your + datastores. + """ + rules: SequenceNotStr[str] """A list of constraints that the agent must follow.""" @@ -48,6 +84,21 @@ class ClientRunAgentParams(TypedDict, total=False): system: str """A system prompt to provide system information to the agent.""" + workflows: Iterable[Workflow] + """A set of declarative workflows for the agent to execute. + + Each workflow is a DAG (Directed Acyclic Graph) of steps that the agent + interprets and executes dynamically. + """ + + +class Callback(TypedDict, total=False): + url: Required[str] + """The webhook URL to send events to.""" + + headers: Dict[str, str] + """Optional headers to include in the callback request.""" + class Datastore(TypedDict, total=False): id: Required[str] @@ -57,9 +108,74 @@ class Datastore(TypedDict, total=False): """The name of the datastore.""" +class FinalOutput(TypedDict, total=False): + name: Required[str] + """Name for the final output.""" + + schema: Required[object] + """JSON schema defining the structure of the final output.""" + + strict: bool + """Whether to enforce strict schema validation.""" + + class McpServer(TypedDict, total=False): server_name: Required[Annotated[str, PropertyInfo(alias="serverName")]] """Name of the MCP server.""" server_url: Required[Annotated[str, PropertyInfo(alias="serverUrl")]] """URL of the MCP server.""" + + +class Query(TypedDict, total=False): + description: Required[str] + """Description of what the query does.""" + + name: Required[str] + """Name of the query action.""" + + query: Required[str] + """The SQL query to execute.""" + + +class WorkflowStepRetryPolicy(TypedDict, total=False): + backoff: str + + max_attempts: int + + +class WorkflowStep(TypedDict, total=False): + id: Required[str] + """Unique identifier for the step.""" + + depends_on: Required[SequenceNotStr[str]] + """Array of step IDs that must complete before this step runs.""" + + description: Required[str] + """What the step should accomplish.""" + + name: Required[str] + """Name of the step.""" + + optional: bool + """Whether the step can be skipped if it fails.""" + + output_schema: object + """JSON schema for expected output validation.""" + + retry_policy: WorkflowStepRetryPolicy + """Retry configuration for the step.""" + + +class Workflow(TypedDict, total=False): + id: Required[str] + """Unique identifier for the workflow.""" + + description: Required[str] + """What the workflow accomplishes.""" + + name: Required[str] + """Name of the workflow.""" + + steps: Required[Iterable[WorkflowStep]] + """Array of step objects.""" diff --git a/src/agentbase/types/clear_message_clear_params.py b/src/agentbase/types/message_clear_params.py similarity index 73% rename from src/agentbase/types/clear_message_clear_params.py rename to src/agentbase/types/message_clear_params.py index b02f2cd..9f29ab8 100644 --- a/src/agentbase/types/clear_message_clear_params.py +++ b/src/agentbase/types/message_clear_params.py @@ -4,9 +4,9 @@ from typing_extensions import Required, TypedDict -__all__ = ["ClearMessageClearParams"] +__all__ = ["MessageClearParams"] -class ClearMessageClearParams(TypedDict, total=False): +class MessageClearParams(TypedDict, total=False): session: Required[str] """The session ID to clear messages from.""" diff --git a/src/agentbase/types/clear_message_clear_response.py b/src/agentbase/types/message_clear_response.py similarity index 79% rename from src/agentbase/types/clear_message_clear_response.py rename to src/agentbase/types/message_clear_response.py index 363e04e..432a1f6 100644 --- a/src/agentbase/types/clear_message_clear_response.py +++ b/src/agentbase/types/message_clear_response.py @@ -4,10 +4,10 @@ from .._models import BaseModel -__all__ = ["ClearMessageClearResponse"] +__all__ = ["MessageClearResponse"] -class ClearMessageClearResponse(BaseModel): +class MessageClearResponse(BaseModel): message: Optional[str] = None """Human‑readable status message.""" diff --git a/src/agentbase/types/get_message_retrieve_params.py b/src/agentbase/types/message_get_params.py similarity index 72% rename from src/agentbase/types/get_message_retrieve_params.py rename to src/agentbase/types/message_get_params.py index 09f031f..391d345 100644 --- a/src/agentbase/types/get_message_retrieve_params.py +++ b/src/agentbase/types/message_get_params.py @@ -4,9 +4,9 @@ from typing_extensions import Required, TypedDict -__all__ = ["GetMessageRetrieveParams"] +__all__ = ["MessageGetParams"] -class GetMessageRetrieveParams(TypedDict, total=False): +class MessageGetParams(TypedDict, total=False): session: Required[str] """The session ID to retrieve messages from.""" diff --git a/src/agentbase/types/get_message_retrieve_response.py b/src/agentbase/types/message_get_response.py similarity index 65% rename from src/agentbase/types/get_message_retrieve_response.py rename to src/agentbase/types/message_get_response.py index 7051487..317e561 100644 --- a/src/agentbase/types/get_message_retrieve_response.py +++ b/src/agentbase/types/message_get_response.py @@ -5,10 +5,10 @@ from .._models import BaseModel -__all__ = ["GetMessageRetrieveResponse", "GetMessageRetrieveResponseItem"] +__all__ = ["MessageGetResponse", "MessageGetResponseItem"] -class GetMessageRetrieveResponseItem(BaseModel): +class MessageGetResponseItem(BaseModel): content: str """The textual content of the message.""" @@ -19,4 +19,4 @@ class GetMessageRetrieveResponseItem(BaseModel): """ -GetMessageRetrieveResponse: TypeAlias = List[GetMessageRetrieveResponseItem] +MessageGetResponse: TypeAlias = List[MessageGetResponseItem] diff --git a/tests/api_resources/test_agent.py b/tests/api_resources/test_agent.py new file mode 100644 index 0000000..5af8a7a --- /dev/null +++ b/tests/api_resources/test_agent.py @@ -0,0 +1,214 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from __future__ import annotations + +import os +from typing import Any, cast + +import pytest + +from agentbase import Agentbase, AsyncAgentbase + +base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") + + +class TestAgent: + parametrize = pytest.mark.parametrize("client", [False, True], indirect=True, ids=["loose", "strict"]) + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + def test_method_run(self, client: Agentbase) -> None: + agent_stream = client.agent.run( + message="message", + ) + agent_stream.response.close() + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + def test_method_run_with_all_params(self, client: Agentbase) -> None: + agent_stream = client.agent.run( + message="message", + session="session", + background=True, + callback={ + "url": "https://example.com", + "headers": {"foo": "string"}, + }, + datastores=[ + { + "id": "id", + "name": "name", + } + ], + final_output={ + "name": "name", + "schema": {}, + "strict": True, + }, + mcp_servers=[ + { + "server_name": "serverName", + "server_url": "https://example.com", + } + ], + mode="flash", + queries=[ + { + "description": "description", + "name": "name", + "query": "query", + } + ], + rules=["string"], + streaming_tokens=True, + system="system", + workflows=[ + { + "id": "id", + "description": "description", + "name": "name", + "steps": [ + { + "id": "id", + "depends_on": ["string"], + "description": "description", + "name": "name", + "optional": True, + "output_schema": {}, + "retry_policy": { + "backoff": "backoff", + "max_attempts": 0, + }, + } + ], + } + ], + ) + agent_stream.response.close() + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + def test_raw_response_run(self, client: Agentbase) -> None: + response = client.agent.with_raw_response.run( + message="message", + ) + + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + stream = response.parse() + stream.close() + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + def test_streaming_response_run(self, client: Agentbase) -> None: + with client.agent.with_streaming_response.run( + message="message", + ) as response: + assert not response.is_closed + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + + stream = response.parse() + stream.close() + + assert cast(Any, response.is_closed) is True + + +class TestAsyncAgent: + parametrize = pytest.mark.parametrize( + "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"] + ) + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + async def test_method_run(self, async_client: AsyncAgentbase) -> None: + agent_stream = await async_client.agent.run( + message="message", + ) + await agent_stream.response.aclose() + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + async def test_method_run_with_all_params(self, async_client: AsyncAgentbase) -> None: + agent_stream = await async_client.agent.run( + message="message", + session="session", + background=True, + callback={ + "url": "https://example.com", + "headers": {"foo": "string"}, + }, + datastores=[ + { + "id": "id", + "name": "name", + } + ], + final_output={ + "name": "name", + "schema": {}, + "strict": True, + }, + mcp_servers=[ + { + "server_name": "serverName", + "server_url": "https://example.com", + } + ], + mode="flash", + queries=[ + { + "description": "description", + "name": "name", + "query": "query", + } + ], + rules=["string"], + streaming_tokens=True, + system="system", + workflows=[ + { + "id": "id", + "description": "description", + "name": "name", + "steps": [ + { + "id": "id", + "depends_on": ["string"], + "description": "description", + "name": "name", + "optional": True, + "output_schema": {}, + "retry_policy": { + "backoff": "backoff", + "max_attempts": 0, + }, + } + ], + } + ], + ) + await agent_stream.response.aclose() + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + async def test_raw_response_run(self, async_client: AsyncAgentbase) -> None: + response = await async_client.agent.with_raw_response.run( + message="message", + ) + + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + stream = await response.parse() + await stream.close() + + @pytest.mark.skip(reason="Prism doesn't support text/event-stream responses") + @parametrize + async def test_streaming_response_run(self, async_client: AsyncAgentbase) -> None: + async with async_client.agent.with_streaming_response.run( + message="message", + ) as response: + assert not response.is_closed + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + + stream = await response.parse() + await stream.close() + + assert cast(Any, response.is_closed) is True diff --git a/tests/api_resources/test_clear_messages.py b/tests/api_resources/test_clear_messages.py deleted file mode 100644 index 09c5846..0000000 --- a/tests/api_resources/test_clear_messages.py +++ /dev/null @@ -1,92 +0,0 @@ -# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. - -from __future__ import annotations - -import os -from typing import Any, cast - -import pytest - -from agentbase import Agentbase, AsyncAgentbase -from tests.utils import assert_matches_type -from agentbase.types import ClearMessageClearResponse - -base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") - - -class TestClearMessages: - parametrize = pytest.mark.parametrize("client", [False, True], indirect=True, ids=["loose", "strict"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - def test_method_clear(self, client: Agentbase) -> None: - clear_message = client.clear_messages.clear( - session="session", - ) - assert_matches_type(ClearMessageClearResponse, clear_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - def test_raw_response_clear(self, client: Agentbase) -> None: - response = client.clear_messages.with_raw_response.clear( - session="session", - ) - - assert response.is_closed is True - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - clear_message = response.parse() - assert_matches_type(ClearMessageClearResponse, clear_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - def test_streaming_response_clear(self, client: Agentbase) -> None: - with client.clear_messages.with_streaming_response.clear( - session="session", - ) as response: - assert not response.is_closed - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - - clear_message = response.parse() - assert_matches_type(ClearMessageClearResponse, clear_message, path=["response"]) - - assert cast(Any, response.is_closed) is True - - -class TestAsyncClearMessages: - parametrize = pytest.mark.parametrize( - "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"] - ) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - async def test_method_clear(self, async_client: AsyncAgentbase) -> None: - clear_message = await async_client.clear_messages.clear( - session="session", - ) - assert_matches_type(ClearMessageClearResponse, clear_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - async def test_raw_response_clear(self, async_client: AsyncAgentbase) -> None: - response = await async_client.clear_messages.with_raw_response.clear( - session="session", - ) - - assert response.is_closed is True - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - clear_message = await response.parse() - assert_matches_type(ClearMessageClearResponse, clear_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - async def test_streaming_response_clear(self, async_client: AsyncAgentbase) -> None: - async with async_client.clear_messages.with_streaming_response.clear( - session="session", - ) as response: - assert not response.is_closed - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - - clear_message = await response.parse() - assert_matches_type(ClearMessageClearResponse, clear_message, path=["response"]) - - assert cast(Any, response.is_closed) is True diff --git a/tests/api_resources/test_client.py b/tests/api_resources/test_client.py index 8a5d15a..3b9ad83 100644 --- a/tests/api_resources/test_client.py +++ b/tests/api_resources/test_client.py @@ -29,12 +29,22 @@ def test_method_run_agent_with_all_params(self, client: Agentbase) -> None: client_stream = client.run_agent( message="message", session="session", + background=True, + callback={ + "url": "https://example.com", + "headers": {"foo": "string"}, + }, datastores=[ { "id": "id", "name": "name", } ], + final_output={ + "name": "name", + "schema": {}, + "strict": True, + }, mcp_servers=[ { "server_name": "serverName", @@ -42,9 +52,37 @@ def test_method_run_agent_with_all_params(self, client: Agentbase) -> None: } ], mode="flash", + queries=[ + { + "description": "description", + "name": "name", + "query": "query", + } + ], rules=["string"], streaming_tokens=True, system="system", + workflows=[ + { + "id": "id", + "description": "description", + "name": "name", + "steps": [ + { + "id": "id", + "depends_on": ["string"], + "description": "description", + "name": "name", + "optional": True, + "output_schema": {}, + "retry_policy": { + "backoff": "backoff", + "max_attempts": 0, + }, + } + ], + } + ], ) client_stream.response.close() @@ -93,12 +131,22 @@ async def test_method_run_agent_with_all_params(self, async_client: AsyncAgentba client_stream = await async_client.run_agent( message="message", session="session", + background=True, + callback={ + "url": "https://example.com", + "headers": {"foo": "string"}, + }, datastores=[ { "id": "id", "name": "name", } ], + final_output={ + "name": "name", + "schema": {}, + "strict": True, + }, mcp_servers=[ { "server_name": "serverName", @@ -106,9 +154,37 @@ async def test_method_run_agent_with_all_params(self, async_client: AsyncAgentba } ], mode="flash", + queries=[ + { + "description": "description", + "name": "name", + "query": "query", + } + ], rules=["string"], streaming_tokens=True, system="system", + workflows=[ + { + "id": "id", + "description": "description", + "name": "name", + "steps": [ + { + "id": "id", + "depends_on": ["string"], + "description": "description", + "name": "name", + "optional": True, + "output_schema": {}, + "retry_policy": { + "backoff": "backoff", + "max_attempts": 0, + }, + } + ], + } + ], ) await client_stream.response.aclose() diff --git a/tests/api_resources/test_get_messages.py b/tests/api_resources/test_get_messages.py deleted file mode 100644 index 884d749..0000000 --- a/tests/api_resources/test_get_messages.py +++ /dev/null @@ -1,92 +0,0 @@ -# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. - -from __future__ import annotations - -import os -from typing import Any, cast - -import pytest - -from agentbase import Agentbase, AsyncAgentbase -from tests.utils import assert_matches_type -from agentbase.types import GetMessageRetrieveResponse - -base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") - - -class TestGetMessages: - parametrize = pytest.mark.parametrize("client", [False, True], indirect=True, ids=["loose", "strict"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - def test_method_retrieve(self, client: Agentbase) -> None: - get_message = client.get_messages.retrieve( - session="session", - ) - assert_matches_type(GetMessageRetrieveResponse, get_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - def test_raw_response_retrieve(self, client: Agentbase) -> None: - response = client.get_messages.with_raw_response.retrieve( - session="session", - ) - - assert response.is_closed is True - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - get_message = response.parse() - assert_matches_type(GetMessageRetrieveResponse, get_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - def test_streaming_response_retrieve(self, client: Agentbase) -> None: - with client.get_messages.with_streaming_response.retrieve( - session="session", - ) as response: - assert not response.is_closed - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - - get_message = response.parse() - assert_matches_type(GetMessageRetrieveResponse, get_message, path=["response"]) - - assert cast(Any, response.is_closed) is True - - -class TestAsyncGetMessages: - parametrize = pytest.mark.parametrize( - "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"] - ) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - async def test_method_retrieve(self, async_client: AsyncAgentbase) -> None: - get_message = await async_client.get_messages.retrieve( - session="session", - ) - assert_matches_type(GetMessageRetrieveResponse, get_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - async def test_raw_response_retrieve(self, async_client: AsyncAgentbase) -> None: - response = await async_client.get_messages.with_raw_response.retrieve( - session="session", - ) - - assert response.is_closed is True - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - get_message = await response.parse() - assert_matches_type(GetMessageRetrieveResponse, get_message, path=["response"]) - - @pytest.mark.skip(reason="Prism tests are disabled") - @parametrize - async def test_streaming_response_retrieve(self, async_client: AsyncAgentbase) -> None: - async with async_client.get_messages.with_streaming_response.retrieve( - session="session", - ) as response: - assert not response.is_closed - assert response.http_request.headers.get("X-Stainless-Lang") == "python" - - get_message = await response.parse() - assert_matches_type(GetMessageRetrieveResponse, get_message, path=["response"]) - - assert cast(Any, response.is_closed) is True diff --git a/tests/api_resources/test_messages.py b/tests/api_resources/test_messages.py new file mode 100644 index 0000000..1ba0209 --- /dev/null +++ b/tests/api_resources/test_messages.py @@ -0,0 +1,160 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from __future__ import annotations + +import os +from typing import Any, cast + +import pytest + +from agentbase import Agentbase, AsyncAgentbase +from tests.utils import assert_matches_type +from agentbase.types import MessageGetResponse, MessageClearResponse + +base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") + + +class TestMessages: + parametrize = pytest.mark.parametrize("client", [False, True], indirect=True, ids=["loose", "strict"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + def test_method_clear(self, client: Agentbase) -> None: + message = client.messages.clear( + session="session", + ) + assert_matches_type(MessageClearResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + def test_raw_response_clear(self, client: Agentbase) -> None: + response = client.messages.with_raw_response.clear( + session="session", + ) + + assert response.is_closed is True + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + message = response.parse() + assert_matches_type(MessageClearResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + def test_streaming_response_clear(self, client: Agentbase) -> None: + with client.messages.with_streaming_response.clear( + session="session", + ) as response: + assert not response.is_closed + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + + message = response.parse() + assert_matches_type(MessageClearResponse, message, path=["response"]) + + assert cast(Any, response.is_closed) is True + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + def test_method_get(self, client: Agentbase) -> None: + message = client.messages.get( + session="session", + ) + assert_matches_type(MessageGetResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + def test_raw_response_get(self, client: Agentbase) -> None: + response = client.messages.with_raw_response.get( + session="session", + ) + + assert response.is_closed is True + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + message = response.parse() + assert_matches_type(MessageGetResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + def test_streaming_response_get(self, client: Agentbase) -> None: + with client.messages.with_streaming_response.get( + session="session", + ) as response: + assert not response.is_closed + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + + message = response.parse() + assert_matches_type(MessageGetResponse, message, path=["response"]) + + assert cast(Any, response.is_closed) is True + + +class TestAsyncMessages: + parametrize = pytest.mark.parametrize( + "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"] + ) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + async def test_method_clear(self, async_client: AsyncAgentbase) -> None: + message = await async_client.messages.clear( + session="session", + ) + assert_matches_type(MessageClearResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + async def test_raw_response_clear(self, async_client: AsyncAgentbase) -> None: + response = await async_client.messages.with_raw_response.clear( + session="session", + ) + + assert response.is_closed is True + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + message = await response.parse() + assert_matches_type(MessageClearResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + async def test_streaming_response_clear(self, async_client: AsyncAgentbase) -> None: + async with async_client.messages.with_streaming_response.clear( + session="session", + ) as response: + assert not response.is_closed + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + + message = await response.parse() + assert_matches_type(MessageClearResponse, message, path=["response"]) + + assert cast(Any, response.is_closed) is True + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + async def test_method_get(self, async_client: AsyncAgentbase) -> None: + message = await async_client.messages.get( + session="session", + ) + assert_matches_type(MessageGetResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + async def test_raw_response_get(self, async_client: AsyncAgentbase) -> None: + response = await async_client.messages.with_raw_response.get( + session="session", + ) + + assert response.is_closed is True + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + message = await response.parse() + assert_matches_type(MessageGetResponse, message, path=["response"]) + + @pytest.mark.skip(reason="Prism tests are disabled") + @parametrize + async def test_streaming_response_get(self, async_client: AsyncAgentbase) -> None: + async with async_client.messages.with_streaming_response.get( + session="session", + ) as response: + assert not response.is_closed + assert response.http_request.headers.get("X-Stainless-Lang") == "python" + + message = await response.parse() + assert_matches_type(MessageGetResponse, message, path=["response"]) + + assert cast(Any, response.is_closed) is True