keepit-ai 0.1.0__py3-none-any.whl

keepit/agents/agent.py

@@ -0,0 +1,191 @@
+"""
+Base Agent interface for Keepit.
+
+This module defines the abstract Agent class, which describes
+the behavior and interface that all agents in Keepit must follow.
+
+Agents can:
+- Provide a textual 'description' of their capabilities,
+- Expose an 'agent_type' to facilitate registry logic,
+- Initialize themselves with 'setup()',
+- Handle incoming messages via 'handle_message()',
+- Dynamically call external tools via 'call_tool()',
+- Discover available tools via 'discover_tools()',
+- Optionally retrieve conversation memory (summary, last n messages)
+  through a MemoryTool if registered in the Tool registry.
+"""
+
+
+import abc
+from typing import Any, Dict, List, Optional
+from dataclasses import dataclass
+from keepit.tools.tool import Tool
+from keepit.tools.tool_registry import ToolRegistry
+from keepit.memory.repository import Repository
+
+@dataclass
+class AgentConfig:
+    """
+    Configuration data for an agent.
+    """
+    agent_name: str
+    agent_type: str
+    description: str
+    system_prompt: str = "You are a helpful AI assistant."
+    llm_config: Optional[Dict[str, Any]] = None
+    tool_registry: Optional[ToolRegistry] = None
+    memory: Optional[Repository] = None
+    is_tool_caller: bool = False
+    is_streaming: bool = False
+
+    def __post_init__(self):
+        if not self.agent_name:
+            raise ValueError("Agent name must be provided.")
+        if not self.description:
+            raise ValueError("Agent description must be provided.")
+        default_llm_config = {
+                'model_name': "default",
+                'temperature': 0.7,
+                'max_tokens':  2000,
+                'top_p': 1.0,
+                'frequency_penalty': 0.0,
+                'presence_penalty': 0.0,
+                'stop_sequences':  [],
+        }
+        self.llm_config = {**default_llm_config, **(self.llm_config or {})}
+
+
+
+class Agent(abc.ABC):
+    """
+    Abstract base class for all Keepit agents.
+
+    Agents are responsible for:
+    - A textual description of their capabilities (description property),
+    - A type descriptor (agent_type) that helps the registry handle them,
+    - Any necessary setup or initialization (setup()),
+    - Receiving messages or prompts (handle_message()),
+    - Generating responses (usually via an LLM or other logic),
+    - Optionally discovering & calling tools (call_tool, discover_tools).
+
+    Concrete implementations (e.g., OpenAIAgent, AnthropicAgent, RemoteAgent)
+    should override the abstract methods with vendor or application-specific logic.
+    """
+
+    def __init__(
+        self,
+        config: AgentConfig
+    ):
+        """
+        Initialize the agent with:
+          - a name (agent_name),
+          - an agent type (agent_type) for registry usage,
+          - a brief description of its capabilities/role,
+          - an optional config dictionary (API keys, model parameters, etc.),
+          - an optional reference to a tool registry.
+
+        :param agent_name: A unique name or identifier for the agent.
+        :param agent_type: A short string representing the type of the agent
+                           (e.g., "BaseAgent", "RemoteAgent", "OpenAIAgent").
+        :param description: A brief explanation of the agent's capabilities and role.
+        :param config: Optional configuration dict (e.g., API keys, parameters).
+        :param agent_config: Optional AgentConfig object with model parameters.
+        :param tool_registry: A reference to a centralized ToolRegistry (if any).
+        """
+        self.agent_name = config.agent_name
+        self.agent_type = config.agent_type
+        self.description = config.description
+        self.llm_config = config.llm_config or {}
+        self.tool_registry = config.tool_registry
+        self.system_prompt = config.system_prompt or "You are a helpful AI assistant."
+        self.memory = config.memory
+        self.is_tool_caller = config.is_tool_caller
+        self.is_streaming = config.is_streaming
+        
+
+    @abc.abstractmethod
+    def handle_message(self, message: str, **kwargs) -> str:
+        """
+        Receive a message (prompt) and generate a response.
+
+        :param message: The user or system prompt to be handled.
+        :param kwargs: Additional context or parameters the agent might need,
+                       such as conversation ID, user metadata, etc.
+        :return: The agent's response as a string.
+        """
+        raise NotImplementedError("Subclasses must implement handle_message().")
+
+    @abc.abstractmethod
+    def handle_message_stream(self, message: str, **kwargs):
+        """
+        Receive a message (prompt) and generate a response in a streaming fashion.
+
+        :param message: The user or system prompt to be handled.
+        :param kwargs: Additional context or parameters the agent might need,
+                       such as conversation ID, user metadata, etc.
+        :yield: Chunks of the agent's response as strings.
+        """
+        raise NotImplementedError("Subclasses must implement handle_message_stream().")
+
+    def call_tool(self, tool_name: str, method_name: str, *args, **kwargs) -> Any:
+        """
+        Call a method on a registered tool by name.
+
+        :param tool_name: The unique name or identifier of the tool.
+        :param method_name: The name of the method to call on the tool.
+        :param args: Positional arguments to pass to the tool method.
+        :param kwargs: Keyword arguments to pass to the tool method.
+        :return: The result of the tool method call.
+        """
+        if not self.tool_registry:
+            raise RuntimeError(
+                f"Agent '{self.agent_name}' has no tool registry attached."
+            )
+
+        tool = self.tool_registry.get_tool(tool_name)
+        if not tool:
+            raise ValueError(
+                f"No tool named '{tool_name}' found in the registry."
+            )
+
+        method = getattr(tool, method_name, None)
+        if not callable(method):
+            raise AttributeError(
+                f"Tool '{tool_name}' does not have method '{method_name}'."
+            )
+
+        return method(*args, **kwargs)
+
+    def discover_tools(self) -> List[str]:
+        """
+        Return a list of available tool names from the registry.
+
+        :return: A list of tool names (strings).
+        """
+        if not self.tool_registry:
+            return []
+        return self.tool_registry.list_tools()
+
+    def get_conversation_summary(self, thread_id: str) -> str:
+        """
+        Retrieve a summary of the conversation so far using the MemoryTool, if available.
+
+        :param thread_id: The identifier of the conversation thread.
+        :return: A textual summary of the conversation so far. If no MemoryTool
+                 or no registry is found, returns an empty string or raises an error.
+        """        
+        if not self.memory:
+            return ""
+        return self.memory.get_conversation_summary(thread_id)
+
+    def get_last_n_messages(self, thread_id: str, n: int = 5) -> List[Any]:
+        """
+        Retrieve the last n messages of the conversation using the MemoryTool, if available.
+
+        :param thread_id: The identifier of the conversation thread.
+        :param n: The number of recent messages to retrieve.
+        :return: A list of message objects or dictionaries.
+        """
+        if not self.memory:
+            return ""
+        return self.memory.get_last_n_messages(thread_id, n)

keepit/agents/agent_info.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class AgentInfo():
+    "A class that holds information about an agent"
+    name: str
+    description: str
+    type: str
+
+    def __init__(self, name: str, description: str, type: str):
+        self.name = name
+        self.description = description
+        self.type = type

keepit/agents/azure_openai_agent.py

@@ -0,0 +1,72 @@
+"""
+AzureOpenAIAgent for Keepit.
+
+An Agent that uses OpenAI's ChatCompletion or Completion API
+to generate responses, pulling API key from the environment.
+"""
+
+
+import os
+from openai import AzureOpenAI
+from dataclasses import dataclass
+
+from typing import Any, Dict, List, Optional
+from keepit.agents.openai_agent import OpenAIAgent, OpenAIAgentConfig
+try :
+    from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+except ImportError:
+    raise ImportError("Azure dependencies for Keepit are not installed. Please install it using 'pip install keepit-ai[azure]'.")
+
+
+
+@dataclass
+class AzureOpenAIAgentConfig(OpenAIAgentConfig):
+    """
+    Configuration data for an AzureOpenAIAgent.
+    """
+    api_base: Optional[str] = None
+    api_version: Optional[str] = None
+    organization: Optional[str] = None
+    use_azure_ad_token_provider: Optional[bool] = False
+
+
+class AzureOpenAIAgent(OpenAIAgent):
+    """
+    A simple AzureOpenAI-based agent that uses the ChatCompletion API.
+    """
+
+    def __init__(self, config: AzureOpenAIAgentConfig):
+        """
+        Initialize the AzureOpenAIAgent.
+
+        :param config: Configuration for the agent.
+        """
+        if config.use_azure_ad_token_provider:
+            config.api_key = "Using Azure AD token provider"
+
+        super().__init__(config=config)
+        if not config.api_base:
+            raise ValueError("Azure OpenAI API base is required for AzureOpenAIAgent.")
+        api_base = config.api_base 
+
+        if not config.api_version:
+            raise ValueError("Azure OpenAI API version is required for AzureOpenAIAgent.")
+        
+
+        
+        if config.use_azure_ad_token_provider:
+            token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+
+            self.client = AzureOpenAI( 
+                                  azure_endpoint=api_base, 
+                                  api_version=config.api_version,
+                                  azure_ad_token_provider=token_provider,
+                                  organization=config.organization)
+        else:
+            self.client = AzureOpenAI(api_key=config.api_key,
+                                  azure_endpoint=api_base,
+                                  api_version=config.api_version,
+                                  organization=config.organization)
+        

keepit/agents/bedrock_agent.py

@@ -0,0 +1,168 @@
+"""
+BedrockAgent for Keepit.
+
+An Agent that uses AWS Bedrock API to generate responses,
+pulling AWS credentials from environment or AWS configuration.
+"""
+
+import json
+import boto3
+from typing import Any, Dict, Optional
+from keepit.agents.agent import Agent, AgentConfig
+from dataclasses import dataclass
+
+
+@dataclass
+class BedrockAgentConfig(AgentConfig):
+    model_id: str = "anthropic.claude-v2"
+    region: str = "us-east-1"
+
+
+
+class BedrockAgent(Agent):
+    """
+    A simple AWS Bedrock-based agent that uses the Bedrock API.
+    """
+
+    def __init__(
+        self,
+        config: BedrockAgentConfig
+    ):
+        """
+        :param agent_name: Unique name or identifier for the agent.
+        :param description: A brief explanation of the agent's capabilities.
+        :param model_id: The Bedrock model ID (e.g., "anthropic.claude-v2").
+        :param config: Optional config dict (can include AWS region).
+        :param tool_registry: Optional ToolRegistry to enable tool calling.
+        :param system_prompt: Default system prompt for context.
+        """
+        super().__init__(
+            config=config
+        )
+        self.config = config
+        self.system_prompt = config.system_prompt
+        self.model_id = config.model_id
+        self.region = config.region
+
+    def setup(self) -> None:
+        """
+        Initialize the Bedrock client using boto3.
+        AWS credentials should be configured via environment variables
+        or AWS configuration files.
+        """
+        try:
+            self.client = boto3.client(
+                service_name='bedrock-runtime',
+                region_name=self.region
+            )
+        except Exception as e:
+            raise EnvironmentError(
+                f"Failed to initialize Bedrock client: {str(e)}"
+            )
+
+    def handle_message(self, message: str, **kwargs) -> str:
+        """
+        Calls AWS Bedrock to handle the user's message.
+        """
+        try:
+            # Construct the request based on the model
+            if "anthropic.claude-3" in self.model_id:
+                # Use Messages API format for Claude 3 models
+                body = {
+                    "anthropic_version": "bedrock-2023-05-31",
+                    "max_tokens": self.config.llm_config.get('max_tokens', 1000),
+                    "temperature": self.config.llm_config.get('temperature', 0.7),
+                    "messages": [
+                        {
+                            "role": "assistant",
+                            "content": self.system_prompt
+                        },
+                        {
+                            "role": "user",
+                            "content": message
+                        }
+                    ]
+                }
+            elif "anthropic" in self.model_id:
+                # Legacy format for older Claude models
+                prompt = f"\n\nHuman: {message}\n\nAssistant:"
+                body = {
+                    "prompt": self.system_prompt + prompt,
+                    "max_tokens_to_sample": self.config.llm_config.get('max_tokens', 1000),
+                    "temperature": self.config.llm_config.get('temperature', 0.7)
+                }
+            else:
+                # Handle other model types here
+                body = {
+                    "inputText": message
+                }
+                
+            response = self.client.invoke_model(
+                modelId=self.model_id,
+                body=json.dumps(body)
+            )
+            response_body = json.loads(response['body'].read())
+            
+            # Handle different response formats
+            if "anthropic.claude-3" in self.model_id:
+                return response_body.get('content', [{}])[0].get('text', '')
+            else:
+                return response_body.get('completion', response_body.get('outputText', ''))
+                
+        except Exception as e:
+            return f"[BedrockAgent error: {str(e)}]"
+            
+    def handle_message_stream(self, message: str, **kwargs):
+        """
+        Calls AWS Bedrock to handle the user's message with streaming support.
+        """
+        try:
+            if "anthropic.claude-3" in self.model_id:
+                # Use Messages API format for Claude 3 models
+                body = {
+                    "anthropic_version": "bedrock-2023-05-31",
+                    "max_tokens": self.config.llm_config.get('max_tokens', 1000),
+                    "temperature": self.config.llm_config.get('temperature', 0.7),
+                    "messages": [
+                        {
+                            "role": "assistant",
+                            "content": self.system_prompt
+                        },
+                        {
+                            "role": "user",
+                            "content": message
+                        }
+                    ]
+                }
+            elif "anthropic" in self.model_id:
+                # Legacy format for older Claude models
+                prompt = f"\n\nHuman: {message}\n\nAssistant:"
+                body = {
+                    "prompt": self.system_prompt + prompt,
+                    "max_tokens_to_sample": self.config.llm_config.get('max_tokens', 1000),
+                    "temperature": self.config.llm_config.get('temperature', 0.7)
+                }
+            else:
+                body = {
+                    "inputText": message
+                }
+                
+            response = self.client.invoke_model_with_response_stream(
+                modelId=self.model_id,
+                body=json.dumps(body)
+            )
+            
+            for event in response['body']:
+                chunk = json.loads(event['chunk']['bytes'])
+                if "anthropic.claude-3" in self.model_id:
+                    if 'delta' in chunk and 'text' in chunk['delta']:
+                        yield chunk['delta']['text']
+                elif 'completion' in chunk:
+                    yield chunk['completion']
+                elif 'outputText' in chunk:
+                    yield chunk['outputText']
+                    
+        except Exception as e:
+            error_message = f"[BedrockAgent error: {str(e)}]"
+            print(error_message)
+            yield error_message

keepit/agents/crewai_agent.py

@@ -0,0 +1,108 @@
+"""
+CrewAIAgent for Keepit.
+
+An Agent that uses a Crew to generate responses using CrewAI.
+"""
+import os
+from dataclasses import dataclass
+
+from crewai import Agent as CrewAgent, LLM as CrewLLM, Task as CrewTask, Crew
+from typing import Any, Dict, Optional
+from keepit.agents.agent import Agent, AgentConfig
+
+os.environ["OTEL_SDK_DISABLED"] = "true"
+
+
+@dataclass
+class CrewAIAgentConfig(AgentConfig):
+    api_key: str = os.getenv("OPENAI_API_KEY"),
+    model_name: str = "gpt-4o"
+
+
+class CrewAIAgent(Agent):
+    """
+    A simple CrewAI-based agent.
+    """
+
+    def __init__(
+            self,
+            agent_name: str,
+            description: str,
+            config: Optional[Dict[str, Any]] = None,
+            tool_registry: Optional[Any] = None,
+            agent_config: Optional[CrewAIAgentConfig] = None
+    ):
+        """
+        :param agent_name: Unique name or identifier for the agent.
+        :param description: A brief explanation of the agent's capabilities.
+        :param config: Optional agent configuration (unused by default).
+        :param tool_registry: Optional ToolRegistry to enable tool calling.
+        :param system_prompt: Default system prompt for context.
+        :param agent_config: Optional configuration for the CrewAIAgent.
+        """
+        super().__init__(
+            agent_name=agent_name,
+            agent_type="CrewAIAgent",
+            description=description,
+            config=config,
+            tool_registry=tool_registry
+        )
+        self.agent_config = agent_config or CrewAIAgentConfig()
+        self.system_prompt = self.agent_config.system_prompt
+        self.client = None
+
+    def setup(self) -> None:
+        """
+        Initialize the CrewAI agent with the provided configuration.
+        """
+        try:
+            self.client = CrewAgent(
+                role="assistant",
+                goal=self.system_prompt,
+                backstory=self.description,
+                verbose=False,
+                llm=CrewLLM(
+                    model=self.agent_config.model_name,
+                    api_key=self.agent_config.api_key,
+                ))
+        except Exception as e:
+            raise EnvironmentError(
+                f"Failed to initialize Crew Agent: {str(e)}"
+            )
+
+    def handle_message(self, message: str, **kwargs) -> str:
+        """
+        Calls the CrewAI agent to handle the user's message.
+        """
+        try:
+            task = CrewTask(
+                description=message,
+                expected_output="",
+                agent=self.client,
+            )
+            crew = Crew(agents=[self.client], tasks=[task])
+            response = crew.kickoff().raw
+            return response
+
+        except Exception as e:
+            return f"[BedrockAgent error: {str(e)}]"
+
+    def handle_message_stream(self, message: str, **kwargs):
+        """
+        Calls the CrewAI agent to handle the user's message.
+        CrewAI does not support streaming responses, so this method is the same as handle_message.
+        """
+        try:
+            task = CrewTask(
+                description=message,
+                expected_output="",
+                agent=self.client,
+            )
+            crew = Crew(agents=[self.client], tasks=[task])
+            response = crew.kickoff().raw
+            yield response
+
+        except Exception as e:
+            error_message = f"[CrewAIAgent error: {str(e)}]"
+            print(error_message)
+            yield error_message