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.
Files changed (46) hide show
  1. keepit_ai-0.1.0/LICENSE +21 -0
  2. keepit_ai-0.1.0/PKG-INFO +88 -0
  3. keepit_ai-0.1.0/README.md +18 -0
  4. keepit_ai-0.1.0/keepit/__init__.py +0 -0
  5. keepit_ai-0.1.0/keepit/agents/__init__.py +0 -0
  6. keepit_ai-0.1.0/keepit/agents/agent.py +191 -0
  7. keepit_ai-0.1.0/keepit/agents/agent_info.py +14 -0
  8. keepit_ai-0.1.0/keepit/agents/azure_openai_agent.py +72 -0
  9. keepit_ai-0.1.0/keepit/agents/bedrock_agent.py +168 -0
  10. keepit_ai-0.1.0/keepit/agents/crewai_agent.py +108 -0
  11. keepit_ai-0.1.0/keepit/agents/gemini_agent.py +393 -0
  12. keepit_ai-0.1.0/keepit/agents/groq_agent.py +56 -0
  13. keepit_ai-0.1.0/keepit/agents/ollama_agent.py +122 -0
  14. keepit_ai-0.1.0/keepit/agents/openai_agent.py +263 -0
  15. keepit_ai-0.1.0/keepit/agents/remote_agent.py +129 -0
  16. keepit_ai-0.1.0/keepit/classifiers/classifier.py +20 -0
  17. keepit_ai-0.1.0/keepit/classifiers/llm_classifier.py +51 -0
  18. keepit_ai-0.1.0/keepit/conversation/__init__.py +0 -0
  19. keepit_ai-0.1.0/keepit/conversation/message.py +72 -0
  20. keepit_ai-0.1.0/keepit/conversation/thread.py +69 -0
  21. keepit_ai-0.1.0/keepit/demo.py +390 -0
  22. keepit_ai-0.1.0/keepit/memory/__init__.py +0 -0
  23. keepit_ai-0.1.0/keepit/memory/file_system_repo.py +172 -0
  24. keepit_ai-0.1.0/keepit/memory/in_memory_repository.py +49 -0
  25. keepit_ai-0.1.0/keepit/memory/repository.py +65 -0
  26. keepit_ai-0.1.0/keepit/orchestrators/__init__.py +0 -0
  27. keepit_ai-0.1.0/keepit/orchestrators/multi_agent_orchestrator.py +88 -0
  28. keepit_ai-0.1.0/keepit/orchestrators/orchestrator.py +55 -0
  29. keepit_ai-0.1.0/keepit/orchestrators/react_orchestrator.py +179 -0
  30. keepit_ai-0.1.0/keepit/orchestrators/simple_orchestrator.py +83 -0
  31. keepit_ai-0.1.0/keepit/registry/__init__.py +0 -0
  32. keepit_ai-0.1.0/keepit/registry/agent_registry.py +86 -0
  33. keepit_ai-0.1.0/keepit/registry/agent_repository.py +56 -0
  34. keepit_ai-0.1.0/keepit/registry/in_memory_agent_repository.py +50 -0
  35. keepit_ai-0.1.0/keepit/tools/__init__.py +0 -0
  36. keepit_ai-0.1.0/keepit/tools/ephemeral_memory.py +111 -0
  37. keepit_ai-0.1.0/keepit/tools/tool.py +155 -0
  38. keepit_ai-0.1.0/keepit/tools/tool_registry.py +165 -0
  39. keepit_ai-0.1.0/keepit/utils/constants.py +9 -0
  40. keepit_ai-0.1.0/keepit_ai.egg-info/PKG-INFO +88 -0
  41. keepit_ai-0.1.0/keepit_ai.egg-info/SOURCES.txt +44 -0
  42. keepit_ai-0.1.0/keepit_ai.egg-info/dependency_links.txt +1 -0
  43. keepit_ai-0.1.0/keepit_ai.egg-info/requires.txt +41 -0
  44. keepit_ai-0.1.0/keepit_ai.egg-info/top_level.txt +1 -0
  45. keepit_ai-0.1.0/pyproject.toml +92 -0
  46. keepit_ai-0.1.0/setup.cfg +4 -0
@@ -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.
@@ -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