Skip to content

Commit 21597a4

Browse files
harel-tussiharel-tussi
authored
feat: agents (#321)
* feat: agents * wip * wip * wip * wip * wip * wip * wip * wip * formatter * wip * wip * wip * Revert "wip" This reverts commit f299a84. * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip
1 parent b17684a commit 21597a4

File tree

19 files changed

+948
-3
lines changed

19 files changed

+948
-3
lines changed

‎README.md‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
- [Installation](#Installation) 💿
2626
- [Usage - Chat Completions](#Usage)
2727
- [Maestro](#Maestro)
28+
- [Agents (Beta)](#Agents-Beta)
2829
- [Conversational RAG (Beta)](#Conversational-RAG-Beta)
2930
- [Older Models Support Usage](#Older-Models-Support-Usage)
3031
- [More Models](#More-Models)
@@ -277,6 +278,61 @@ For a more detailed example, see maestro [sync](examples/studio/maestro/run.py)
277278

278279
---
279280

281+
### Agents (Beta)
282+
283+
AI21 Agents provide a comprehensive way to create, manage, and run your Agents.
284+
285+
```python
286+
from ai21 import AI21Client
287+
from ai21.models.agents import BudgetLevel, AgentType
288+
289+
client = AI21Client()
290+
291+
# Run the agent
292+
run_response = client.beta.agents.runs.create_and_poll(
293+
agent_id=agent.id,
294+
input=[{"role": "user", "content": "What is 2+2?"}],
295+
poll_timeout_sec=120,
296+
)
297+
298+
print(f"Result: {run_response.result}")
299+
300+
```
301+
302+
#### Agent CRUD Operations
303+
304+
```python
305+
from ai21 import AI21Client
306+
from ai21.models.agents import BudgetLevel, AgentType
307+
308+
client = AI21Client()
309+
310+
# Create
311+
agent = client.beta.agents.create(
312+
name="Research Assistant",
313+
description="Specialized in research tasks",
314+
budget=BudgetLevel.HIGH,
315+
)
316+
317+
# Read
318+
retrieved_agent = client.beta.agents.get(agent.id)
319+
agents_list = client.beta.agents.list()
320+
321+
# Update
322+
modified_agent = client.beta.agents.modify(
323+
agent.id,
324+
name="Enhanced Research Assistant",
325+
description="Updated with enhanced capabilities",
326+
)
327+
328+
# Delete
329+
delete_response = client.beta.agents.delete(agent.id)
330+
```
331+
332+
For more detailed examples, see agent [CRUD operations](examples/studio/agents/agent_crud.py), [basic runs](examples/studio/agents/agent_run.py), and [async operations](examples/studio/agents/async_agent_run.py) examples.
333+
334+
---
335+
280336
### Conversational RAG (Beta)
281337

282338
Like chat, but with the ability to retrieve information from your Studio library.

‎ai21/clients/common/agents/__init__.py‎

Whitespace-only changes.

‎ai21/clients/common/agents/agents.py‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
from __future__ import annotations
2+
3+
from abc import ABC, abstractmethod
4+
from typing import List, Dict, Any
5+
6+
from ai21.models.agents import (
7+
Agent,
8+
BudgetLevel,
9+
DeleteAgentResponse,
10+
ListAgentsResponse,
11+
AgentRequirement,
12+
Visibility,
13+
)
14+
from ai21.models.agents.agent import ResponseLanguage
15+
from ai21.types import NOT_GIVEN, NotGiven
16+
from ai21.utils.typing import remove_not_given
17+
18+
19+
class BaseAgents(ABC):
20+
_module_name = "assistants"
21+
22+
def _create_body(
23+
self,
24+
*,
25+
name: str,
26+
description: str | NotGiven,
27+
models: List[str] | NotGiven,
28+
tools: List[Dict[str, Any]] | NotGiven,
29+
tool_resources: Dict[str, Any] | NotGiven,
30+
requirements: List[AgentRequirement] | NotGiven,
31+
budget: BudgetLevel | NotGiven,
32+
response_language: ResponseLanguage | NotGiven,
33+
**kwargs,
34+
) -> dict:
35+
return remove_not_given(
36+
{
37+
"name": name,
38+
"description": description,
39+
"models": models,
40+
"tools": tools,
41+
"tool_resources": tool_resources,
42+
"requirements": requirements,
43+
"budget": budget,
44+
"response_language": response_language,
45+
**kwargs,
46+
}
47+
)
48+
49+
def _modify_body(
50+
self,
51+
*,
52+
name: str | NotGiven,
53+
description: str | NotGiven,
54+
models: List[str] | NotGiven,
55+
tools: List[Dict[str, Any]] | NotGiven,
56+
tool_resources: Dict[str, Any] | NotGiven,
57+
requirements: List[AgentRequirement] | NotGiven,
58+
budget: BudgetLevel | NotGiven,
59+
visibility: Visibility | NotGiven,
60+
response_language: ResponseLanguage | NotGiven,
61+
**kwargs,
62+
) -> dict:
63+
return remove_not_given(
64+
{
65+
"name": name,
66+
"description": description,
67+
"models": models,
68+
"tools": tools,
69+
"tool_resources": tool_resources,
70+
"requirements": requirements,
71+
"budget": budget,
72+
"visibility": visibility,
73+
"response_language": response_language,
74+
**kwargs,
75+
}
76+
)
77+
78+
@abstractmethod
79+
def create(
80+
self,
81+
*,
82+
name: str,
83+
description: str | NotGiven = NOT_GIVEN,
84+
optimization: str | NotGiven = NOT_GIVEN,
85+
avatar: str | NotGiven = NOT_GIVEN,
86+
models: List[str] | NotGiven = NOT_GIVEN,
87+
tools: List[Dict[str, Any]] | NotGiven = NOT_GIVEN,
88+
tool_resources: Dict[str, Any] | NotGiven = NOT_GIVEN,
89+
requirements: List[AgentRequirement] | NotGiven = NOT_GIVEN,
90+
budget: BudgetLevel | NotGiven = NOT_GIVEN,
91+
response_language: ResponseLanguage | NotGiven = NOT_GIVEN,
92+
**kwargs,
93+
) -> Agent:
94+
pass
95+
96+
@abstractmethod
97+
def get(self, agent_id: str) -> Agent:
98+
pass
99+
100+
@abstractmethod
101+
def list(self) -> ListAgentsResponse:
102+
pass
103+
104+
@abstractmethod
105+
def modify(
106+
self,
107+
agent_id: str,
108+
*,
109+
name: str | NotGiven = NOT_GIVEN,
110+
description: str | NotGiven = NOT_GIVEN,
111+
models: List[str] | NotGiven = NOT_GIVEN,
112+
tools: List[Dict[str, Any]] | NotGiven = NOT_GIVEN,
113+
tool_resources: Dict[str, Any] | NotGiven = NOT_GIVEN,
114+
requirements: List[AgentRequirement] | NotGiven = NOT_GIVEN,
115+
budget: BudgetLevel | NotGiven = NOT_GIVEN,
116+
visibility: Visibility | NotGiven = NOT_GIVEN,
117+
response_language: ResponseLanguage | NotGiven = NOT_GIVEN,
118+
**kwargs,
119+
) -> Agent:
120+
pass
121+
122+
@abstractmethod
123+
def delete(self, agent_id: str) -> DeleteAgentResponse:
124+
pass

‎ai21/clients/common/agents/run.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from abc import ABC
2+
from typing import List
3+
4+
from ai21.models.agents import Agent
5+
from ai21.models.agents.agent import AgentRequirement
6+
from ai21.models.maestro.run import Requirement
7+
from ai21.utils.typing import remove_not_given
8+
9+
10+
class BaseAgentRun(ABC):
11+
def _convert_requirements(self, requirements: List[AgentRequirement]):
12+
"""Convert agent requirements to maestro requirements, filtering out invalid ones."""
13+
if not requirements:
14+
return None
15+
16+
converted_requirements = []
17+
for req in requirements:
18+
if req.title and req.description and req.type:
19+
converted_requirements.append(
20+
Requirement(
21+
name=req.title,
22+
description=req.description,
23+
is_mandatory=req.type == "mandatory",
24+
)
25+
)
26+
27+
return converted_requirements or None
28+
29+
def convert_agent_to_maestro_run_payload(
30+
self,
31+
agent: Agent,
32+
**kwargs,
33+
):
34+
return remove_not_given(
35+
{
36+
"models": agent.models,
37+
"tools": agent.tools,
38+
"tool_resources": agent.tool_resources,
39+
"requirements": self._convert_requirements(agent.requirements),
40+
"budget": agent.budget,
41+
"response_language": agent.response_language,
42+
**kwargs,
43+
}
44+
)

‎ai21/clients/common/maestro/maestro.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from abc import ABC
2-
32
from ai21.clients.common.maestro.run import BaseMaestroRun
43

54

‎ai21/clients/studio/resources/agents/__init__.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
from ai21.clients.studio.resources.agents.agents import Agents
2+
from ai21.clients.studio.resources.agents.async_agents import AsyncAgents
3+
4+
__all__ = [
5+
"Agents",
6+
"AsyncAgents",
7+
]

0 commit comments

Comments
 (0)