keepit-ai 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- keepit_ai-0.1.0/LICENSE +21 -0
- keepit_ai-0.1.0/PKG-INFO +88 -0
- keepit_ai-0.1.0/README.md +18 -0
- keepit_ai-0.1.0/keepit/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/agents/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/agents/agent.py +191 -0
- keepit_ai-0.1.0/keepit/agents/agent_info.py +14 -0
- keepit_ai-0.1.0/keepit/agents/azure_openai_agent.py +72 -0
- keepit_ai-0.1.0/keepit/agents/bedrock_agent.py +168 -0
- keepit_ai-0.1.0/keepit/agents/crewai_agent.py +108 -0
- keepit_ai-0.1.0/keepit/agents/gemini_agent.py +393 -0
- keepit_ai-0.1.0/keepit/agents/groq_agent.py +56 -0
- keepit_ai-0.1.0/keepit/agents/ollama_agent.py +122 -0
- keepit_ai-0.1.0/keepit/agents/openai_agent.py +263 -0
- keepit_ai-0.1.0/keepit/agents/remote_agent.py +129 -0
- keepit_ai-0.1.0/keepit/classifiers/classifier.py +20 -0
- keepit_ai-0.1.0/keepit/classifiers/llm_classifier.py +51 -0
- keepit_ai-0.1.0/keepit/conversation/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/conversation/message.py +72 -0
- keepit_ai-0.1.0/keepit/conversation/thread.py +69 -0
- keepit_ai-0.1.0/keepit/demo.py +390 -0
- keepit_ai-0.1.0/keepit/memory/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/memory/file_system_repo.py +172 -0
- keepit_ai-0.1.0/keepit/memory/in_memory_repository.py +49 -0
- keepit_ai-0.1.0/keepit/memory/repository.py +65 -0
- keepit_ai-0.1.0/keepit/orchestrators/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/orchestrators/multi_agent_orchestrator.py +88 -0
- keepit_ai-0.1.0/keepit/orchestrators/orchestrator.py +55 -0
- keepit_ai-0.1.0/keepit/orchestrators/react_orchestrator.py +179 -0
- keepit_ai-0.1.0/keepit/orchestrators/simple_orchestrator.py +83 -0
- keepit_ai-0.1.0/keepit/registry/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/registry/agent_registry.py +86 -0
- keepit_ai-0.1.0/keepit/registry/agent_repository.py +56 -0
- keepit_ai-0.1.0/keepit/registry/in_memory_agent_repository.py +50 -0
- keepit_ai-0.1.0/keepit/tools/__init__.py +0 -0
- keepit_ai-0.1.0/keepit/tools/ephemeral_memory.py +111 -0
- keepit_ai-0.1.0/keepit/tools/tool.py +155 -0
- keepit_ai-0.1.0/keepit/tools/tool_registry.py +165 -0
- keepit_ai-0.1.0/keepit/utils/constants.py +9 -0
- keepit_ai-0.1.0/keepit_ai.egg-info/PKG-INFO +88 -0
- keepit_ai-0.1.0/keepit_ai.egg-info/SOURCES.txt +44 -0
- keepit_ai-0.1.0/keepit_ai.egg-info/dependency_links.txt +1 -0
- keepit_ai-0.1.0/keepit_ai.egg-info/requires.txt +41 -0
- keepit_ai-0.1.0/keepit_ai.egg-info/top_level.txt +1 -0
- keepit_ai-0.1.0/pyproject.toml +92 -0
- keepit_ai-0.1.0/setup.cfg +4 -0
keepit_ai-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Developer
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
keepit_ai-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: keepit-ai
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A multi-agent orchestration framework with advanced memory management and Gemini/MCP tool calling.
|
|
5
|
+
Author-email: Developer <dev@example.com>
|
|
6
|
+
License: MIT License
|
|
7
|
+
|
|
8
|
+
Copyright (c) 2026 Developer
|
|
9
|
+
|
|
10
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
11
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
12
|
+
in the Software without restriction, including without limitation the rights
|
|
13
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
14
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
15
|
+
furnished to do so, subject to the following conditions:
|
|
16
|
+
|
|
17
|
+
The above copyright notice and this permission notice shall be included in all
|
|
18
|
+
copies or substantial portions of the Software.
|
|
19
|
+
|
|
20
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
21
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
22
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
23
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
24
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
25
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
26
|
+
SOFTWARE.
|
|
27
|
+
|
|
28
|
+
Keywords: python,keepit,agents,mcp,memory
|
|
29
|
+
Classifier: Development Status :: 3 - Alpha
|
|
30
|
+
Classifier: Programming Language :: Python :: 3
|
|
31
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
32
|
+
Classifier: Operating System :: OS Independent
|
|
33
|
+
Requires-Python: >=3.11
|
|
34
|
+
Description-Content-Type: text/markdown
|
|
35
|
+
License-File: LICENSE
|
|
36
|
+
Requires-Dist: mcp>=0.1.0
|
|
37
|
+
Requires-Dist: nest-asyncio>=1.6.0
|
|
38
|
+
Requires-Dist: python-dotenv>=1.0.1
|
|
39
|
+
Provides-Extra: openai
|
|
40
|
+
Requires-Dist: openai>=1.60.1; extra == "openai"
|
|
41
|
+
Provides-Extra: awsbedrock
|
|
42
|
+
Requires-Dist: boto3>=1.36.9; extra == "awsbedrock"
|
|
43
|
+
Provides-Extra: crewai
|
|
44
|
+
Requires-Dist: crewai>=0.100.1; extra == "crewai"
|
|
45
|
+
Requires-Dist: crewai-tools>=0.33.0; extra == "crewai"
|
|
46
|
+
Provides-Extra: ollama
|
|
47
|
+
Requires-Dist: requests>=2.32.3; extra == "ollama"
|
|
48
|
+
Provides-Extra: gemini
|
|
49
|
+
Requires-Dist: google-genai>=0.1.1; extra == "gemini"
|
|
50
|
+
Provides-Extra: groq
|
|
51
|
+
Requires-Dist: groq>=0.11.0; extra == "groq"
|
|
52
|
+
Provides-Extra: azure
|
|
53
|
+
Requires-Dist: openai>=1.60.1; extra == "azure"
|
|
54
|
+
Requires-Dist: azure-identity>=1.21.0; extra == "azure"
|
|
55
|
+
Provides-Extra: all
|
|
56
|
+
Requires-Dist: boto3>=1.36.9; extra == "all"
|
|
57
|
+
Requires-Dist: crewai>=0.100.1; extra == "all"
|
|
58
|
+
Requires-Dist: crewai-tools>=0.33.0; extra == "all"
|
|
59
|
+
Requires-Dist: requests>=2.32.3; extra == "all"
|
|
60
|
+
Requires-Dist: fastapi>=0.115.7; extra == "all"
|
|
61
|
+
Requires-Dist: uvicorn>=0.34.0; extra == "all"
|
|
62
|
+
Requires-Dist: python-dotenv>=1.0.1; extra == "all"
|
|
63
|
+
Requires-Dist: openai>=1.60.1; extra == "all"
|
|
64
|
+
Requires-Dist: azure-identity>=1.21.0; extra == "all"
|
|
65
|
+
Requires-Dist: mcp>=0.1.0; extra == "all"
|
|
66
|
+
Requires-Dist: nest-asyncio>=1.6.0; extra == "all"
|
|
67
|
+
Requires-Dist: google-genai>=0.1.1; extra == "all"
|
|
68
|
+
Requires-Dist: groq>=0.11.0; extra == "all"
|
|
69
|
+
Dynamic: license-file
|
|
70
|
+
|
|
71
|
+
# Keepit - Meta Orchestration framework for Your Agents
|
|
72
|
+
|
|
73
|
+
Keepit is a multi-agent framework that coordinates messages between users and specialized AI agents. It uses routing orchestrators, classifiers, and tool registries to delegate tasks to agents, managing conversation state via a modular memory subsystem. It also implements Model Context Protocol (MCP) tool integration.
|
|
74
|
+
|
|
75
|
+
## Installation
|
|
76
|
+
|
|
77
|
+
```bash
|
|
78
|
+
pip install .
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
## Features
|
|
82
|
+
|
|
83
|
+
- **Agent Management**: Create, register, and manage multiple AI agents.
|
|
84
|
+
- **Orchestration**: Orchestrate conversations and tasks across multiple agents.
|
|
85
|
+
- **Memory Tools**: Integrate memory tools to maintain conversation context and history.
|
|
86
|
+
- **MCP Client**: Dynamically load tools from MCP Servers over Stdio or SSE transport.
|
|
87
|
+
- **Streaming Responses**: Support for streaming responses from agents.
|
|
88
|
+
- **Extensibility**: Easily extend the framework with new agents, tools, and orchestrators.
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
# Keepit - Meta Orchestration framework for Your Agents
|
|
2
|
+
|
|
3
|
+
Keepit is a multi-agent framework that coordinates messages between users and specialized AI agents. It uses routing orchestrators, classifiers, and tool registries to delegate tasks to agents, managing conversation state via a modular memory subsystem. It also implements Model Context Protocol (MCP) tool integration.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install .
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Features
|
|
12
|
+
|
|
13
|
+
- **Agent Management**: Create, register, and manage multiple AI agents.
|
|
14
|
+
- **Orchestration**: Orchestrate conversations and tasks across multiple agents.
|
|
15
|
+
- **Memory Tools**: Integrate memory tools to maintain conversation context and history.
|
|
16
|
+
- **MCP Client**: Dynamically load tools from MCP Servers over Stdio or SSE transport.
|
|
17
|
+
- **Streaming Responses**: Support for streaming responses from agents.
|
|
18
|
+
- **Extensibility**: Easily extend the framework with new agents, tools, and orchestrators.
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,191 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Base Agent interface for Keepit.
|
|
3
|
+
|
|
4
|
+
This module defines the abstract Agent class, which describes
|
|
5
|
+
the behavior and interface that all agents in Keepit must follow.
|
|
6
|
+
|
|
7
|
+
Agents can:
|
|
8
|
+
- Provide a textual 'description' of their capabilities,
|
|
9
|
+
- Expose an 'agent_type' to facilitate registry logic,
|
|
10
|
+
- Initialize themselves with 'setup()',
|
|
11
|
+
- Handle incoming messages via 'handle_message()',
|
|
12
|
+
- Dynamically call external tools via 'call_tool()',
|
|
13
|
+
- Discover available tools via 'discover_tools()',
|
|
14
|
+
- Optionally retrieve conversation memory (summary, last n messages)
|
|
15
|
+
through a MemoryTool if registered in the Tool registry.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
import abc
|
|
20
|
+
from typing import Any, Dict, List, Optional
|
|
21
|
+
from dataclasses import dataclass
|
|
22
|
+
from keepit.tools.tool import Tool
|
|
23
|
+
from keepit.tools.tool_registry import ToolRegistry
|
|
24
|
+
from keepit.memory.repository import Repository
|
|
25
|
+
|
|
26
|
+
@dataclass
|
|
27
|
+
class AgentConfig:
|
|
28
|
+
"""
|
|
29
|
+
Configuration data for an agent.
|
|
30
|
+
"""
|
|
31
|
+
agent_name: str
|
|
32
|
+
agent_type: str
|
|
33
|
+
description: str
|
|
34
|
+
system_prompt: str = "You are a helpful AI assistant."
|
|
35
|
+
llm_config: Optional[Dict[str, Any]] = None
|
|
36
|
+
tool_registry: Optional[ToolRegistry] = None
|
|
37
|
+
memory: Optional[Repository] = None
|
|
38
|
+
is_tool_caller: bool = False
|
|
39
|
+
is_streaming: bool = False
|
|
40
|
+
|
|
41
|
+
def __post_init__(self):
|
|
42
|
+
if not self.agent_name:
|
|
43
|
+
raise ValueError("Agent name must be provided.")
|
|
44
|
+
if not self.description:
|
|
45
|
+
raise ValueError("Agent description must be provided.")
|
|
46
|
+
default_llm_config = {
|
|
47
|
+
'model_name': "default",
|
|
48
|
+
'temperature': 0.7,
|
|
49
|
+
'max_tokens': 2000,
|
|
50
|
+
'top_p': 1.0,
|
|
51
|
+
'frequency_penalty': 0.0,
|
|
52
|
+
'presence_penalty': 0.0,
|
|
53
|
+
'stop_sequences': [],
|
|
54
|
+
}
|
|
55
|
+
self.llm_config = {**default_llm_config, **(self.llm_config or {})}
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class Agent(abc.ABC):
|
|
60
|
+
"""
|
|
61
|
+
Abstract base class for all Keepit agents.
|
|
62
|
+
|
|
63
|
+
Agents are responsible for:
|
|
64
|
+
- A textual description of their capabilities (description property),
|
|
65
|
+
- A type descriptor (agent_type) that helps the registry handle them,
|
|
66
|
+
- Any necessary setup or initialization (setup()),
|
|
67
|
+
- Receiving messages or prompts (handle_message()),
|
|
68
|
+
- Generating responses (usually via an LLM or other logic),
|
|
69
|
+
- Optionally discovering & calling tools (call_tool, discover_tools).
|
|
70
|
+
|
|
71
|
+
Concrete implementations (e.g., OpenAIAgent, AnthropicAgent, RemoteAgent)
|
|
72
|
+
should override the abstract methods with vendor or application-specific logic.
|
|
73
|
+
"""
|
|
74
|
+
|
|
75
|
+
def __init__(
|
|
76
|
+
self,
|
|
77
|
+
config: AgentConfig
|
|
78
|
+
):
|
|
79
|
+
"""
|
|
80
|
+
Initialize the agent with:
|
|
81
|
+
- a name (agent_name),
|
|
82
|
+
- an agent type (agent_type) for registry usage,
|
|
83
|
+
- a brief description of its capabilities/role,
|
|
84
|
+
- an optional config dictionary (API keys, model parameters, etc.),
|
|
85
|
+
- an optional reference to a tool registry.
|
|
86
|
+
|
|
87
|
+
:param agent_name: A unique name or identifier for the agent.
|
|
88
|
+
:param agent_type: A short string representing the type of the agent
|
|
89
|
+
(e.g., "BaseAgent", "RemoteAgent", "OpenAIAgent").
|
|
90
|
+
:param description: A brief explanation of the agent's capabilities and role.
|
|
91
|
+
:param config: Optional configuration dict (e.g., API keys, parameters).
|
|
92
|
+
:param agent_config: Optional AgentConfig object with model parameters.
|
|
93
|
+
:param tool_registry: A reference to a centralized ToolRegistry (if any).
|
|
94
|
+
"""
|
|
95
|
+
self.agent_name = config.agent_name
|
|
96
|
+
self.agent_type = config.agent_type
|
|
97
|
+
self.description = config.description
|
|
98
|
+
self.llm_config = config.llm_config or {}
|
|
99
|
+
self.tool_registry = config.tool_registry
|
|
100
|
+
self.system_prompt = config.system_prompt or "You are a helpful AI assistant."
|
|
101
|
+
self.memory = config.memory
|
|
102
|
+
self.is_tool_caller = config.is_tool_caller
|
|
103
|
+
self.is_streaming = config.is_streaming
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
@abc.abstractmethod
|
|
107
|
+
def handle_message(self, message: str, **kwargs) -> str:
|
|
108
|
+
"""
|
|
109
|
+
Receive a message (prompt) and generate a response.
|
|
110
|
+
|
|
111
|
+
:param message: The user or system prompt to be handled.
|
|
112
|
+
:param kwargs: Additional context or parameters the agent might need,
|
|
113
|
+
such as conversation ID, user metadata, etc.
|
|
114
|
+
:return: The agent's response as a string.
|
|
115
|
+
"""
|
|
116
|
+
raise NotImplementedError("Subclasses must implement handle_message().")
|
|
117
|
+
|
|
118
|
+
@abc.abstractmethod
|
|
119
|
+
def handle_message_stream(self, message: str, **kwargs):
|
|
120
|
+
"""
|
|
121
|
+
Receive a message (prompt) and generate a response in a streaming fashion.
|
|
122
|
+
|
|
123
|
+
:param message: The user or system prompt to be handled.
|
|
124
|
+
:param kwargs: Additional context or parameters the agent might need,
|
|
125
|
+
such as conversation ID, user metadata, etc.
|
|
126
|
+
:yield: Chunks of the agent's response as strings.
|
|
127
|
+
"""
|
|
128
|
+
raise NotImplementedError("Subclasses must implement handle_message_stream().")
|
|
129
|
+
|
|
130
|
+
def call_tool(self, tool_name: str, method_name: str, *args, **kwargs) -> Any:
|
|
131
|
+
"""
|
|
132
|
+
Call a method on a registered tool by name.
|
|
133
|
+
|
|
134
|
+
:param tool_name: The unique name or identifier of the tool.
|
|
135
|
+
:param method_name: The name of the method to call on the tool.
|
|
136
|
+
:param args: Positional arguments to pass to the tool method.
|
|
137
|
+
:param kwargs: Keyword arguments to pass to the tool method.
|
|
138
|
+
:return: The result of the tool method call.
|
|
139
|
+
"""
|
|
140
|
+
if not self.tool_registry:
|
|
141
|
+
raise RuntimeError(
|
|
142
|
+
f"Agent '{self.agent_name}' has no tool registry attached."
|
|
143
|
+
)
|
|
144
|
+
|
|
145
|
+
tool = self.tool_registry.get_tool(tool_name)
|
|
146
|
+
if not tool:
|
|
147
|
+
raise ValueError(
|
|
148
|
+
f"No tool named '{tool_name}' found in the registry."
|
|
149
|
+
)
|
|
150
|
+
|
|
151
|
+
method = getattr(tool, method_name, None)
|
|
152
|
+
if not callable(method):
|
|
153
|
+
raise AttributeError(
|
|
154
|
+
f"Tool '{tool_name}' does not have method '{method_name}'."
|
|
155
|
+
)
|
|
156
|
+
|
|
157
|
+
return method(*args, **kwargs)
|
|
158
|
+
|
|
159
|
+
def discover_tools(self) -> List[str]:
|
|
160
|
+
"""
|
|
161
|
+
Return a list of available tool names from the registry.
|
|
162
|
+
|
|
163
|
+
:return: A list of tool names (strings).
|
|
164
|
+
"""
|
|
165
|
+
if not self.tool_registry:
|
|
166
|
+
return []
|
|
167
|
+
return self.tool_registry.list_tools()
|
|
168
|
+
|
|
169
|
+
def get_conversation_summary(self, thread_id: str) -> str:
|
|
170
|
+
"""
|
|
171
|
+
Retrieve a summary of the conversation so far using the MemoryTool, if available.
|
|
172
|
+
|
|
173
|
+
:param thread_id: The identifier of the conversation thread.
|
|
174
|
+
:return: A textual summary of the conversation so far. If no MemoryTool
|
|
175
|
+
or no registry is found, returns an empty string or raises an error.
|
|
176
|
+
"""
|
|
177
|
+
if not self.memory:
|
|
178
|
+
return ""
|
|
179
|
+
return self.memory.get_conversation_summary(thread_id)
|
|
180
|
+
|
|
181
|
+
def get_last_n_messages(self, thread_id: str, n: int = 5) -> List[Any]:
|
|
182
|
+
"""
|
|
183
|
+
Retrieve the last n messages of the conversation using the MemoryTool, if available.
|
|
184
|
+
|
|
185
|
+
:param thread_id: The identifier of the conversation thread.
|
|
186
|
+
:param n: The number of recent messages to retrieve.
|
|
187
|
+
:return: A list of message objects or dictionaries.
|
|
188
|
+
"""
|
|
189
|
+
if not self.memory:
|
|
190
|
+
return ""
|
|
191
|
+
return self.memory.get_last_n_messages(thread_id, n)
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
@dataclass
|
|
5
|
+
class AgentInfo():
|
|
6
|
+
"A class that holds information about an agent"
|
|
7
|
+
name: str
|
|
8
|
+
description: str
|
|
9
|
+
type: str
|
|
10
|
+
|
|
11
|
+
def __init__(self, name: str, description: str, type: str):
|
|
12
|
+
self.name = name
|
|
13
|
+
self.description = description
|
|
14
|
+
self.type = type
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
"""
|
|
2
|
+
AzureOpenAIAgent for Keepit.
|
|
3
|
+
|
|
4
|
+
An Agent that uses OpenAI's ChatCompletion or Completion API
|
|
5
|
+
to generate responses, pulling API key from the environment.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
import os
|
|
10
|
+
from openai import AzureOpenAI
|
|
11
|
+
from dataclasses import dataclass
|
|
12
|
+
|
|
13
|
+
from typing import Any, Dict, List, Optional
|
|
14
|
+
from keepit.agents.openai_agent import OpenAIAgent, OpenAIAgentConfig
|
|
15
|
+
try :
|
|
16
|
+
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
|
|
17
|
+
except ImportError:
|
|
18
|
+
raise ImportError("Azure dependencies for Keepit are not installed. Please install it using 'pip install keepit-ai[azure]'.")
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@dataclass
|
|
23
|
+
class AzureOpenAIAgentConfig(OpenAIAgentConfig):
|
|
24
|
+
"""
|
|
25
|
+
Configuration data for an AzureOpenAIAgent.
|
|
26
|
+
"""
|
|
27
|
+
api_base: Optional[str] = None
|
|
28
|
+
api_version: Optional[str] = None
|
|
29
|
+
organization: Optional[str] = None
|
|
30
|
+
use_azure_ad_token_provider: Optional[bool] = False
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class AzureOpenAIAgent(OpenAIAgent):
|
|
34
|
+
"""
|
|
35
|
+
A simple AzureOpenAI-based agent that uses the ChatCompletion API.
|
|
36
|
+
"""
|
|
37
|
+
|
|
38
|
+
def __init__(self, config: AzureOpenAIAgentConfig):
|
|
39
|
+
"""
|
|
40
|
+
Initialize the AzureOpenAIAgent.
|
|
41
|
+
|
|
42
|
+
:param config: Configuration for the agent.
|
|
43
|
+
"""
|
|
44
|
+
if config.use_azure_ad_token_provider:
|
|
45
|
+
config.api_key = "Using Azure AD token provider"
|
|
46
|
+
|
|
47
|
+
super().__init__(config=config)
|
|
48
|
+
if not config.api_base:
|
|
49
|
+
raise ValueError("Azure OpenAI API base is required for AzureOpenAIAgent.")
|
|
50
|
+
api_base = config.api_base
|
|
51
|
+
|
|
52
|
+
if not config.api_version:
|
|
53
|
+
raise ValueError("Azure OpenAI API version is required for AzureOpenAIAgent.")
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
if config.use_azure_ad_token_provider:
|
|
58
|
+
token_provider = get_bearer_token_provider(
|
|
59
|
+
DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
self.client = AzureOpenAI(
|
|
63
|
+
azure_endpoint=api_base,
|
|
64
|
+
api_version=config.api_version,
|
|
65
|
+
azure_ad_token_provider=token_provider,
|
|
66
|
+
organization=config.organization)
|
|
67
|
+
else:
|
|
68
|
+
self.client = AzureOpenAI(api_key=config.api_key,
|
|
69
|
+
azure_endpoint=api_base,
|
|
70
|
+
api_version=config.api_version,
|
|
71
|
+
organization=config.organization)
|
|
72
|
+
|
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
"""
|
|
2
|
+
BedrockAgent for Keepit.
|
|
3
|
+
|
|
4
|
+
An Agent that uses AWS Bedrock API to generate responses,
|
|
5
|
+
pulling AWS credentials from environment or AWS configuration.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import json
|
|
9
|
+
import boto3
|
|
10
|
+
from typing import Any, Dict, Optional
|
|
11
|
+
from keepit.agents.agent import Agent, AgentConfig
|
|
12
|
+
from dataclasses import dataclass
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@dataclass
|
|
16
|
+
class BedrockAgentConfig(AgentConfig):
|
|
17
|
+
model_id: str = "anthropic.claude-v2"
|
|
18
|
+
region: str = "us-east-1"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class BedrockAgent(Agent):
|
|
23
|
+
"""
|
|
24
|
+
A simple AWS Bedrock-based agent that uses the Bedrock API.
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
def __init__(
|
|
28
|
+
self,
|
|
29
|
+
config: BedrockAgentConfig
|
|
30
|
+
):
|
|
31
|
+
"""
|
|
32
|
+
:param agent_name: Unique name or identifier for the agent.
|
|
33
|
+
:param description: A brief explanation of the agent's capabilities.
|
|
34
|
+
:param model_id: The Bedrock model ID (e.g., "anthropic.claude-v2").
|
|
35
|
+
:param config: Optional config dict (can include AWS region).
|
|
36
|
+
:param tool_registry: Optional ToolRegistry to enable tool calling.
|
|
37
|
+
:param system_prompt: Default system prompt for context.
|
|
38
|
+
"""
|
|
39
|
+
super().__init__(
|
|
40
|
+
config=config
|
|
41
|
+
)
|
|
42
|
+
self.config = config
|
|
43
|
+
self.system_prompt = config.system_prompt
|
|
44
|
+
self.model_id = config.model_id
|
|
45
|
+
self.region = config.region
|
|
46
|
+
|
|
47
|
+
def setup(self) -> None:
|
|
48
|
+
"""
|
|
49
|
+
Initialize the Bedrock client using boto3.
|
|
50
|
+
AWS credentials should be configured via environment variables
|
|
51
|
+
or AWS configuration files.
|
|
52
|
+
"""
|
|
53
|
+
try:
|
|
54
|
+
self.client = boto3.client(
|
|
55
|
+
service_name='bedrock-runtime',
|
|
56
|
+
region_name=self.region
|
|
57
|
+
)
|
|
58
|
+
except Exception as e:
|
|
59
|
+
raise EnvironmentError(
|
|
60
|
+
f"Failed to initialize Bedrock client: {str(e)}"
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
def handle_message(self, message: str, **kwargs) -> str:
|
|
64
|
+
"""
|
|
65
|
+
Calls AWS Bedrock to handle the user's message.
|
|
66
|
+
"""
|
|
67
|
+
try:
|
|
68
|
+
# Construct the request based on the model
|
|
69
|
+
if "anthropic.claude-3" in self.model_id:
|
|
70
|
+
# Use Messages API format for Claude 3 models
|
|
71
|
+
body = {
|
|
72
|
+
"anthropic_version": "bedrock-2023-05-31",
|
|
73
|
+
"max_tokens": self.config.llm_config.get('max_tokens', 1000),
|
|
74
|
+
"temperature": self.config.llm_config.get('temperature', 0.7),
|
|
75
|
+
"messages": [
|
|
76
|
+
{
|
|
77
|
+
"role": "assistant",
|
|
78
|
+
"content": self.system_prompt
|
|
79
|
+
},
|
|
80
|
+
{
|
|
81
|
+
"role": "user",
|
|
82
|
+
"content": message
|
|
83
|
+
}
|
|
84
|
+
]
|
|
85
|
+
}
|
|
86
|
+
elif "anthropic" in self.model_id:
|
|
87
|
+
# Legacy format for older Claude models
|
|
88
|
+
prompt = f"\n\nHuman: {message}\n\nAssistant:"
|
|
89
|
+
body = {
|
|
90
|
+
"prompt": self.system_prompt + prompt,
|
|
91
|
+
"max_tokens_to_sample": self.config.llm_config.get('max_tokens', 1000),
|
|
92
|
+
"temperature": self.config.llm_config.get('temperature', 0.7)
|
|
93
|
+
}
|
|
94
|
+
else:
|
|
95
|
+
# Handle other model types here
|
|
96
|
+
body = {
|
|
97
|
+
"inputText": message
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
response = self.client.invoke_model(
|
|
101
|
+
modelId=self.model_id,
|
|
102
|
+
body=json.dumps(body)
|
|
103
|
+
)
|
|
104
|
+
response_body = json.loads(response['body'].read())
|
|
105
|
+
|
|
106
|
+
# Handle different response formats
|
|
107
|
+
if "anthropic.claude-3" in self.model_id:
|
|
108
|
+
return response_body.get('content', [{}])[0].get('text', '')
|
|
109
|
+
else:
|
|
110
|
+
return response_body.get('completion', response_body.get('outputText', ''))
|
|
111
|
+
|
|
112
|
+
except Exception as e:
|
|
113
|
+
return f"[BedrockAgent error: {str(e)}]"
|
|
114
|
+
|
|
115
|
+
def handle_message_stream(self, message: str, **kwargs):
|
|
116
|
+
"""
|
|
117
|
+
Calls AWS Bedrock to handle the user's message with streaming support.
|
|
118
|
+
"""
|
|
119
|
+
try:
|
|
120
|
+
if "anthropic.claude-3" in self.model_id:
|
|
121
|
+
# Use Messages API format for Claude 3 models
|
|
122
|
+
body = {
|
|
123
|
+
"anthropic_version": "bedrock-2023-05-31",
|
|
124
|
+
"max_tokens": self.config.llm_config.get('max_tokens', 1000),
|
|
125
|
+
"temperature": self.config.llm_config.get('temperature', 0.7),
|
|
126
|
+
"messages": [
|
|
127
|
+
{
|
|
128
|
+
"role": "assistant",
|
|
129
|
+
"content": self.system_prompt
|
|
130
|
+
},
|
|
131
|
+
{
|
|
132
|
+
"role": "user",
|
|
133
|
+
"content": message
|
|
134
|
+
}
|
|
135
|
+
]
|
|
136
|
+
}
|
|
137
|
+
elif "anthropic" in self.model_id:
|
|
138
|
+
# Legacy format for older Claude models
|
|
139
|
+
prompt = f"\n\nHuman: {message}\n\nAssistant:"
|
|
140
|
+
body = {
|
|
141
|
+
"prompt": self.system_prompt + prompt,
|
|
142
|
+
"max_tokens_to_sample": self.config.llm_config.get('max_tokens', 1000),
|
|
143
|
+
"temperature": self.config.llm_config.get('temperature', 0.7)
|
|
144
|
+
}
|
|
145
|
+
else:
|
|
146
|
+
body = {
|
|
147
|
+
"inputText": message
|
|
148
|
+
}
|
|
149
|
+
|
|
150
|
+
response = self.client.invoke_model_with_response_stream(
|
|
151
|
+
modelId=self.model_id,
|
|
152
|
+
body=json.dumps(body)
|
|
153
|
+
)
|
|
154
|
+
|
|
155
|
+
for event in response['body']:
|
|
156
|
+
chunk = json.loads(event['chunk']['bytes'])
|
|
157
|
+
if "anthropic.claude-3" in self.model_id:
|
|
158
|
+
if 'delta' in chunk and 'text' in chunk['delta']:
|
|
159
|
+
yield chunk['delta']['text']
|
|
160
|
+
elif 'completion' in chunk:
|
|
161
|
+
yield chunk['completion']
|
|
162
|
+
elif 'outputText' in chunk:
|
|
163
|
+
yield chunk['outputText']
|
|
164
|
+
|
|
165
|
+
except Exception as e:
|
|
166
|
+
error_message = f"[BedrockAgent error: {str(e)}]"
|
|
167
|
+
print(error_message)
|
|
168
|
+
yield error_message
|