From 45e0eb5175c5ae0f1dbd48d487dc1c6203208075 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 28 Oct 2024 08:41:10 -0400 Subject: [PATCH 1/5] chore: Add base classes for chains and agents --- README.md | 13 ++++ simplemind/agents/__init__.py | 0 simplemind/agents/base.py | 7 ++ simplemind/chains/__init__.py | 0 simplemind/chains/base.py | 7 ++ simplemind/chains/reverse_text.py | 6 ++ simplemind/client.py | 35 ++++++++++ simplemind/concepts.py | 18 +++-- simplemind/config.py | 13 ++++ simplemind/core.py | 21 ------ simplemind/integrations/anthropic.py | 70 +++++++++++++------ simplemind/integrations/base.py | 83 ++++++++++++++++------ simplemind/integrations/openai.py | 95 +++++++++++++------------- simplemind/logger.py | 15 ++++ simplemind/models.py | 30 +++++++- simplemind/plugins/__init__.py | 0 simplemind/plugins/basic_memory.py | 10 +++ simplemind/plugins/kv.py | 10 +++ simplemind/vector_store/__init__.py | 0 simplemind/vector_store/faiss_store.py | 19 ++++++ t.py | 68 +++++++----------- t2.py | 72 ++++++++++++++----- t3.py | 15 ++++ tests/test_openai.py | 26 +++++++ 24 files changed, 454 insertions(+), 179 deletions(-) create mode 100644 simplemind/agents/__init__.py create mode 100644 simplemind/agents/base.py create mode 100644 simplemind/chains/__init__.py create mode 100644 simplemind/chains/base.py create mode 100644 simplemind/chains/reverse_text.py create mode 100644 simplemind/client.py create mode 100644 simplemind/config.py create mode 100644 simplemind/logger.py create mode 100644 simplemind/plugins/__init__.py create mode 100644 simplemind/plugins/basic_memory.py create mode 100644 simplemind/plugins/kv.py create mode 100644 simplemind/vector_store/__init__.py create mode 100644 simplemind/vector_store/faiss_store.py create mode 100644 t3.py create mode 100644 tests/test_openai.py diff --git a/README.md b/README.md index bde1855..e639f9d 100644 --- a/README.md +++ b/README.md @@ -104,3 +104,16 @@ SimpleMind is inspired by the philosophy of "code for humans" and aims to make w --- SimpleMind: Keep it simple, keep it human. + +------------------------ + + +## Plugins + + +SimpleMind supports a plugin system to extend its functionality. Currently available plugins: + +- **KVPlugin**: Key-Value storage for context management. +- **BasicMemoryPlugin**: Simple memory storage for conversations. + +**Adding a Plugin:** diff --git a/simplemind/agents/__init__.py b/simplemind/agents/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simplemind/agents/base.py b/simplemind/agents/base.py new file mode 100644 index 0000000..cdc32a2 --- /dev/null +++ b/simplemind/agents/base.py @@ -0,0 +1,7 @@ +from abc import ABC, abstractmethod + + +class BaseAgent(ABC): + @abstractmethod + def decide(self, context, *args, **kwargs): + pass diff --git a/simplemind/chains/__init__.py b/simplemind/chains/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simplemind/chains/base.py b/simplemind/chains/base.py new file mode 100644 index 0000000..279fdf5 --- /dev/null +++ b/simplemind/chains/base.py @@ -0,0 +1,7 @@ +from abc import ABC, abstractmethod + + +class BaseChain(ABC): + @abstractmethod + def run(self, input_data): + pass diff --git a/simplemind/chains/reverse_text.py b/simplemind/chains/reverse_text.py new file mode 100644 index 0000000..7f102a2 --- /dev/null +++ b/simplemind/chains/reverse_text.py @@ -0,0 +1,6 @@ +from .base import BaseChain + + +class ReverseTextChain(BaseChain): + def run(self, input_data): + return input_data[::-1] diff --git a/simplemind/client.py b/simplemind/client.py new file mode 100644 index 0000000..5bb1dbe --- /dev/null +++ b/simplemind/client.py @@ -0,0 +1,35 @@ +from typing import Optional +from simplemind.models import Conversation, AIResponse +from simplemind.concepts import Context +from simplemind.integrations.openai import OpenAI +from simplemind.integrations.anthropic import Anthropic + +class Client: + def __init__(self, api_key: str, context: Optional[Context] = None): + self.api_key = api_key + self.context = context or Context() + self.providers = self._initialize_providers() + + def _initialize_providers(self): + return { + "openai": OpenAI(api_key=self.api_key), + "anthropic": Anthropic(api_key=self.api_key), + } + + def create_conversation(self, provider: str = "openai") -> Conversation: + if provider not in self.providers: + raise ValueError(f"Provider '{provider}' not supported.") + return self.providers[provider].create_conversation(initial_message="Hello!", context=self.context.dict()) + + def send_message(self, conversation: Conversation, message: str, provider: str = "openai") -> AIResponse: + if provider not in self.providers: + raise ValueError(f"Provider '{provider}' not supported.") + return self.providers[provider].send_message(conversation.id, message) + + @property + def available_models(self): + available = {} + for name, provider in self.providers.items(): + available[name] = provider.available_models + return available + diff --git a/simplemind/concepts.py b/simplemind/concepts.py index bd84ae0..147d6a0 100644 --- a/simplemind/concepts.py +++ b/simplemind/concepts.py @@ -1,6 +1,16 @@ -class Context: - def __init__(self): - self.plugins = [kv, basic_memory] +from pydantic import BaseModel +from typing import Dict, Any +from simplemind.plugins.base import BasePlugin -# TODO: explore pluggy for this. +class Context(BaseModel): + plugins: Dict[str, BasePlugin] = {} + + def add_plugin(self, name: str, plugin: BasePlugin): + self.plugins[name] = plugin + + def execute_plugin(self, name: str, *args, **kwargs): + if name in self.plugins: + return self.plugins[name].execute(self, *args, **kwargs) + else: + raise ValueError(f"Plugin '{name}' not found in context.") diff --git a/simplemind/config.py b/simplemind/config.py new file mode 100644 index 0000000..a14bb25 --- /dev/null +++ b/simplemind/config.py @@ -0,0 +1,13 @@ +from pydantic import BaseSettings + + +class Settings(BaseSettings): + openai_api_key: str = "" + anthropic_api_key: str = "" + default_model: str = "gpt-4" + + class Config: + env_file = ".env" + + +settings = Settings() diff --git a/simplemind/core.py b/simplemind/core.py index e45f25f..e69de29 100644 --- a/simplemind/core.py +++ b/simplemind/core.py @@ -1,21 +0,0 @@ -from fastapi import FastAPI, HTTPException -from pydantic import BaseModel -from typing import Any - -app = FastAPI(title="SimpleMind AI API", description="AI for humans, replacing LangGraph and LangChain for Python users.") - - - -@app.post("/generate", response_model=AIResponse) -def generate_response(request: AIRequest): - try: - # Placeholder for AI generation logic - response = {"message": "This would be the AI response."} - metadata = {"tokens_used": 50} - return AIResponse(response=response, metadata=metadata) - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) - -@app.get("/health") -def health_check(): - return {"status": "healthy"} diff --git a/simplemind/integrations/anthropic.py b/simplemind/integrations/anthropic.py index b96a033..15678dc 100644 --- a/simplemind/integrations/anthropic.py +++ b/simplemind/integrations/anthropic.py @@ -1,41 +1,69 @@ import os +from typing import List, Optional import instructor from anthropic import Anthropic as BaseAnthropic from .base import BaseClientProvider +from ..models import AIResponse, Conversation +from ..logger import logger class Anthropic(BaseClientProvider): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, model: str = DEFAULT_MODEL, api_key: Optional[str] = None): + super().__init__(model=model, api_key=api_key) self.login() def login(self): - """Initialize Anthropic client, with Instructor enabled.""" - - # Default to environment variable if not provided. - if self._api_key is None: + if not self._api_key: self._api_key = os.getenv("ANTHROPIC_API_KEY") - + if not self._api_key: + raise ValueError("Anthropic API key not provided.") base_client = BaseAnthropic(api_key=self._api_key) self.client = instructor.from_anthropic(base_client) - # assert self.test_connection() + if not self.test_connection(): + raise ConnectionError("Failed to connect to Anthropic API.") + logger.info("Logged in to Anthropic successfully.") @property - def available_models(self): - """Returns the available models from the Anthropic client.""" + def available_models(self) -> List[str]: + try: + return [ + "claude-3-opus-20240229", + "claude-3-5-sonnet-20240620", + "claude-3-haiku-20240307", + ] + except Exception as e: + logger.error(f"Error fetching models: {e}") + return [] - # TODO: scrape from website or embed - return [ - "claude-3-opus-20240229", - "claude-3-5-sonnet-20240620", - "claude-3-haiku-20240307", - "claude-3-5-sonnet-20240620", - "claude-3-5-sonnet-20240620", + def test_connection(self) -> bool: + models = self.available_models + if models: + logger.info(f"Available models: {models}") + return True + logger.warning("No available models found.") + return False + + def generate_response(self, conversation: Conversation) -> AIResponse: + messages = [ + {"role": msg.role, "content": msg.content} for msg in conversation.messages ] + params = { + "messages": messages, + "model": self.model, + } + if conversation.context: + params["context"] = conversation.context - # def test_connection(self): - # """Test the connection to Anthropic. Returns True if successful.""" - - # raise NotImplementedError("Anthropic test_connection not implemented.") + try: + completion = self.client.completions.create(**params) + response_text = completion.completion + metadata = {"model": completion.model, "usage": completion.usage} + logger.info("Generated response from Anthropic.") + return AIResponse( + text=response_text, response=completion, metadata=metadata + ) + except Exception as e: + logger.error(f"Error generating response: {e}") + raise e diff --git a/simplemind/integrations/base.py b/simplemind/integrations/base.py index 9ca5d3b..361adf9 100644 --- a/simplemind/integrations/base.py +++ b/simplemind/integrations/base.py @@ -1,6 +1,10 @@ # import logging from pydantic import BaseModel +from typing import Any, Dict, List, Optional +from ..models import AIResponse, Conversation, Message +import uuid +from abc import ABC, abstractmethod DEFAULT_MODEL = "gpt-4o" @@ -8,55 +12,46 @@ DEFAULT_MODEL = "gpt-4o" class BaseClientProvider: - def __init__(self, *, model=DEFAULT_MODEL, api_key=None): + def __init__(self, *, model: str = DEFAULT_MODEL, api_key: Optional[str] = None): # self.logger = logging.getLogger(self.__class__.__name__) self.client = None self.model = model - - # Load API key from environment if not provided self._api_key = api_key + self.conversations: Dict[str, Conversation] = {} + @abstractmethod def login(self): """Initializes the AI provider client.""" - msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - def test_connection(self): + @abstractmethod + def test_connection(self) -> bool: """Tests the connection to the AI provider client.""" - msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - # def generate_response(self, request): - # """Generates a response from the AI provider client.""" - - # msg = "This method must be implemented by the AI provider client." - # raise NotImplementedError(msg) - + @abstractmethod def health_check(self): """Checks the health of the AI provider client.""" - msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - @property - def available_models(self): + @abstractmethod + def available_models(self) -> List[str]: """Returns the available models from the AI provider client.""" - msg = "This method must be implemented by the AI provider client." - raise NotImplementedError(msg) - def message(self, message, **kwargs): + @abstractmethod + def message(self, message: str, **kwargs) -> AIResponse: """Generates a response from the AI provider client.""" - msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) + # Uncomment and implement additional methods as needed # def features(self): # """Returns the features of the AI provider client.""" - # msg = "This method must be implemented by the AI provider client." # raise NotImplementedError(msg) @@ -71,3 +66,51 @@ class BaseClientProvider: # def start_conversation(self, model, message, **kwargs): # pass + + def create_conversation( + self, initial_message: str, context: Optional[Dict[str, Any]] = None + ) -> Conversation: + conv_id = str(uuid.uuid4()) + conversation = Conversation( + id=conv_id, + messages=[Message(role="user", content=initial_message)], + context=context or {}, + ) + self.conversations[conv_id] = conversation + return conversation + + def send_message( + self, + conversation_id: str, + message: str, + context_update: Optional[Dict[str, Any]] = None, + ) -> AIResponse: + if conversation_id not in self.conversations: + raise ValueError("Conversation ID does not exist.") + + conversation = self.conversations[conversation_id] + conversation.messages.append(Message(role="user", content=message)) + + if context_update: + conversation.context.update(context_update) + + response = self.generate_response(conversation) + conversation.messages.append(Message(role="assistant", content=response.text)) + return response + + def generate_response(self, conversation: Conversation) -> AIResponse: + """Generates a response based on the conversation.""" + raise NotImplementedError( + "This method must be implemented by the AI provider client." + ) + + def get_conversation(self, conversation_id: str) -> Conversation: + if conversation_id not in self.conversations: + raise ValueError("Conversation ID does not exist.") + return self.conversations[conversation_id] + + +class BasePlugin(ABC): + @abstractmethod + def execute(self, context, *args, **kwargs): + pass diff --git a/simplemind/integrations/openai.py b/simplemind/integrations/openai.py index d120c7e..d67d9f5 100644 --- a/simplemind/integrations/openai.py +++ b/simplemind/integrations/openai.py @@ -1,68 +1,67 @@ import os - +from typing import Optional, List import instructor from openai import OpenAI as BaseOpenAI from .base import BaseClientProvider -from ..models import AIResponse - +from ..models import AIResponse, Conversation +from ..logger import logger +from simplemind.config import settings class OpenAI(BaseClientProvider): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, model: str = DEFAULT_MODEL, api_key: Optional[str] = None): + super().__init__(model=model, api_key=api_key) self.login() def login(self): - """Initialize OpenAI client, with Instructor enabled.""" - - # Default to environment variable if not provided. - if self._api_key is None: - self._api_key = os.getenv("OPENAI_API_KEY") - + if not self._api_key: + self._api_key = settings.openai_api_key + if not self._api_key: + raise ValueError("OpenAI API key not provided.") self.client = BaseOpenAI(api_key=self._api_key) self.instructor_client = instructor.from_openai(self.client) - assert self.test_connection() + if not self.test_connection(): + raise ConnectionError("Failed to connect to OpenAI API.") + logger.info("Logged in to OpenAI successfully.") @property - def available_models(self): - """Returns the available models from the OpenAI client.""" + def available_models(self) -> List[str]: + try: + return [model.id for model in self.client.models.list()] + except Exception as e: + logger.error(f"Error fetching models: {e}") + return [] - def gen(): - for model in self.client.models.list(): - yield model.id + def test_connection(self) -> bool: + try: + models = self.available_models + if models: + logger.info(f"Available models: {models}") + return True + else: + logger.warning("No available models found.") + return False + except Exception as e: + logger.error(f"Error testing connection: {e}") + return False - return [g for g in gen()] - - def test_connection(self): - """Test the connection to OpenAI. Returns True if successful.""" - - return bool(len(self.available_models)) - - def message(self, message, *, response_model=False, **kwargs): - """Generates a response from the OpenAI client.""" - use_instructor = bool(response_model) - - client = self.instructor_client if use_instructor else self.client - - # Parameters for the OpenAI client. + def generate_response(self, conversation: Conversation) -> AIResponse: + messages = [ + {"role": msg.role, "content": msg.content} for msg in conversation.messages + ] params = { - "messages": [{"role": "user", "content": message}], + "messages": messages, "model": self.model, } - params.update(kwargs) + if conversation.context: + params["context"] = conversation.context - if use_instructor: - params["response_model"] = response_model - - # Make the request to OpenAI. - completion = client.chat.completions.create(**params) - - if use_instructor: - return completion.model_dump() - - else: - return AIResponse( - response=completion, - text=completion.choices[0].message.content, - ) + try: + completion = self.client.chat.completions.create(**params) + response_text = completion.choices[0].message.content + metadata = {"model": completion.model, "usage": completion.usage} + logger.info("Generated response from OpenAI.") + return AIResponse(text=response_text, response=completion, metadata=metadata) + except Exception as e: + logger.error(f"Error generating response: {e}") + raise e diff --git a/simplemind/logger.py b/simplemind/logger.py new file mode 100644 index 0000000..8c2e2e7 --- /dev/null +++ b/simplemind/logger.py @@ -0,0 +1,15 @@ +import logging + +def setup_logger(name: str) -> logging.Logger: + logger = logging.getLogger(name) + if not logger.hasHandlers(): + logger.setLevel(logging.INFO) + handler = logging.StreamHandler() + formatter = logging.Formatter('[%(asctime)s] %(levelname)s - %(message)s') + handler.setFormatter(formatter) + logger.addHandler(handler) + return logger + +# Initialize a global logger +logger = setup_logger("simplemind") + diff --git a/simplemind/models.py b/simplemind/models.py index 6caa2f2..ea097e5 100644 --- a/simplemind/models.py +++ b/simplemind/models.py @@ -1,10 +1,11 @@ from pydantic import BaseModel -from typing import Any, ClassVar +from typing import Any, Dict, List, Optional +import uuid class AIRequest(BaseModel): text: str - parameters: dict = {} + parameters: Dict[str, Any] = {} def __str__(self): return self.text @@ -13,7 +14,30 @@ class AIRequest(BaseModel): class AIResponse(BaseModel): text: str response: Any - metadata: dict = {} + metadata: Dict[str, Any] = {} def __str__(self): return self.text + + +class Message(BaseModel): + role: str # "user" or "assistant" + content: str + + +class Conversation(BaseModel): + id: str + messages: List[Message] = [] + context: Optional[Dict[str, Any]] = {} + + +class ConversationRequest(BaseModel): + conversation_id: Optional[str] = None + message: str + context_update: Optional[Dict[str, Any]] = None + + +class ConversationResponse(BaseModel): + conversation_id: str + messages: List[Message] + metadata: Dict[str, Any] = {} diff --git a/simplemind/plugins/__init__.py b/simplemind/plugins/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simplemind/plugins/basic_memory.py b/simplemind/plugins/basic_memory.py new file mode 100644 index 0000000..d3365e4 --- /dev/null +++ b/simplemind/plugins/basic_memory.py @@ -0,0 +1,10 @@ +from .base import BasePlugin + + +class BasicMemoryPlugin(BasePlugin): + def __init__(self): + self.memory = [] + + def execute(self, context, message): + self.memory.append(message) + return self.memory diff --git a/simplemind/plugins/kv.py b/simplemind/plugins/kv.py new file mode 100644 index 0000000..136ca58 --- /dev/null +++ b/simplemind/plugins/kv.py @@ -0,0 +1,10 @@ +from .base import BasePlugin + + +class KVPlugin(BasePlugin): + def __init__(self): + self.store = {} + + def execute(self, context, key, value): + self.store[key] = value + return self.store diff --git a/simplemind/vector_store/__init__.py b/simplemind/vector_store/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simplemind/vector_store/faiss_store.py b/simplemind/vector_store/faiss_store.py new file mode 100644 index 0000000..9bbc559 --- /dev/null +++ b/simplemind/vector_store/faiss_store.py @@ -0,0 +1,19 @@ +import faiss +import numpy as np +from typing import List + + +class FAISSStore: + def __init__(self, dimension: int): + self.dimension = dimension + self.index = faiss.IndexFlatL2(dimension) + self.ids = [] + + def add_embeddings(self, embeddings: np.ndarray, ids: List[str]): + self.index.add(embeddings) + self.ids.extend(ids) + + def search(self, query_embedding: np.ndarray, top_k: int = 5): + distances, indices = self.index.search(query_embedding, top_k) + results = [(self.ids[idx], distances[i]) for i, idx in enumerate(indices[0])] + return results diff --git a/t.py b/t.py index de740d8..ff34f6b 100644 --- a/t.py +++ b/t.py @@ -1,54 +1,34 @@ from pprint import pprint from pydantic import BaseModel import simplemind - -context = None - -openai = simplemind.integrations.OpenAI() +from simplemind.concepts import Context +from simplemind.plugins.kv import KVPlugin +from simplemind.plugins.basic_memory import BasicMemoryPlugin +from simplemind.chains.reverse_text import ReverseTextChain +from simplemind.client import Client -class YearlyData(BaseModel): - year: int - events: list[str] +class MyContext(Context): + def __init__(self): + super().__init__() + self.add_plugin("kv", KVPlugin()) + self.add_plugin("basic_memory", BasicMemoryPlugin()) -class ProjectData(BaseModel): - name: str - description: str - url: str - github_url: str +# Initialize context and client +context = MyContext() +aiclient = Client(api_key="YOUR_API_KEY", context=context) +# Test connection and available models +print(aiclient.available_models) -class BioData(BaseModel): - bio: str - spouse_name: str - history: list[YearlyData] - fun_facts: list[str] - # age: int - # occupation: str - # bio: str - # affiliations: list[str] +# Example usage +conversation = aiclient.create_conversation(provider="openai") +response = aiclient.send_message( + conversation, "Who is Kenneth Reitz?", provider="openai" +) +print(response) - -class PersonData(BaseModel): - bio: BioData - projects: list[ProjectData] - yearly_breakdown: list[YearlyData] - - -print(openai.test_connection()) -print(openai.available_models) - -print() -print() -message = "who is kenneth reitz?" - -print(f"> {message}") -pprint(openai.message(message, response_model=BioData)) - -# claude = simplemind.integrations.Anthropic() - -# # print(claude.test_connection()) -# # print(claude.available_models) - -# claude.login() +reverse_chain = ReverseTextChain() +result = reverse_chain.run("Hello, World!") +print(result) # Output: !dlroW ,olleH diff --git a/t2.py b/t2.py index 81c4627..0bed67f 100644 --- a/t2.py +++ b/t2.py @@ -1,32 +1,68 @@ -import instructor +from pprint import pprint from pydantic import BaseModel -from openai import OpenAI +import simplemind +from simplemind.vector_store.faiss_store import FAISSStore +import numpy as np + +context = None + +openai = simplemind.integrations.OpenAI() -class ProjectInfo(BaseModel): +class YearlyData(BaseModel): + year: int + events: list[str] + + +class ProjectData(BaseModel): name: str description: str url: str github_url: str -# Define your desired output structure -class UserInfo(BaseModel): - name: str - age: int +class BioData(BaseModel): bio: str - projects: list[ProjectInfo] + spouse_name: str + history: list[YearlyData] + fun_facts: list[str] + # age: int + # occupation: str + # bio: str + # affiliations: list[str] -# Patch the OpenAI client -client = instructor.from_openai(OpenAI()) +class PersonData(BaseModel): + bio: BioData + projects: list[ProjectData] + yearly_breakdown: list[YearlyData] -# Extract structured data from natural language -user_info = client.chat.completions.create( - model="gpt-4o", - response_model=UserInfo, - messages=[{"role": "user", "content": "who is kennethreitz?"}], -) -print(user_info.model_dump()) -# > 30 +print(openai.test_connection()) +print(openai.available_models) + +print() +print() +message = "who is kenneth reitz?" + +print(f"> {message}") +pprint(openai.message(message, response_model=BioData)) + +# claude = simplemind.integrations.Anthropic() + +# # print(claude.test_connection()) +# # print(claude.available_models) + +# claude.login() + +vector_store = FAISSStore(dimension=768) # Example dimension for embeddings + +# Add embeddings +embeddings = np.random.random((10, 768)).astype('float32') +ids = [f"doc_{i}" for i in range(10)] +vector_store.add_embeddings(embeddings, ids) + +# Search +query_embedding = np.random.random((1, 768)).astype('float32') +results = vector_store.search(query_embedding, top_k=3) +print(results) diff --git a/t3.py b/t3.py new file mode 100644 index 0000000..fbe6905 --- /dev/null +++ b/t3.py @@ -0,0 +1,15 @@ +from simplemind.concepts import Context + +# from simplemind.plugins.default_plugin import DefaultPlugin + +# Initialize the context +ctx = Context() + +# Add and initialize the DefaultPlugin +# ctx.add_plugin(DefaultPlugin, "DefaultPlugin") + +# Execute the DefaultPlugin with some data +# ctx.execute_plugin("DefaultPlugin", {"key": "value"}) + +# Shutdown all plugins +# ctx.shutdown_plugins() diff --git a/tests/test_openai.py b/tests/test_openai.py new file mode 100644 index 0000000..8b08252 --- /dev/null +++ b/tests/test_openai.py @@ -0,0 +1,26 @@ +import unittest +from unittest.mock import patch, MagicMock +from simplemind.integrations.openai import OpenAI + + +class TestOpenAIProvider(unittest.TestCase): + @patch("simplemind.integrations.openai.BaseOpenAI") + def setUp(self, mock_openai): + self.mock_openai = mock_openai.return_value + self.mock_openai.models.list.return_value = [MagicMock(id="gpt-4")] + self.provider = OpenAI(api_key="test_api_key", model="gpt-4") + + def test_available_models(self): + models = self.provider.available_models + self.assertIn("gpt-4", models) + + def test_test_connection_success(self): + self.assertTrue(self.provider.test_connection()) + + def test_generate_response_not_implemented(self): + with self.assertRaises(NotImplementedError): + self.provider.generate_response(None) + + +if __name__ == "__main__": + unittest.main() From aaf3f942690567c5106012153e68a9813f5760aa Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 28 Oct 2024 09:00:59 -0400 Subject: [PATCH 2/5] test --- Pipfile | 11 +++++++ simplemind/client.py | 23 ++++++++++++--- simplemind/concepts.py | 1 + simplemind/config.py | 2 +- simplemind/core.py | 44 ++++++++++++++++++++++++++++ simplemind/integrations/anthropic.py | 3 ++ simplemind/integrations/base.py | 17 +++-------- simplemind/integrations/openai.py | 31 ++++++++++++-------- simplemind/models.py | 18 ++++++++++-- simplemind/plugins/base.py | 29 ++++++++++++++++++ simplemind/plugins/kv.py | 14 +++++++-- t.py | 12 +++++--- 12 files changed, 166 insertions(+), 39 deletions(-) create mode 100644 Pipfile create mode 100644 simplemind/plugins/base.py diff --git a/Pipfile b/Pipfile new file mode 100644 index 0000000..0757494 --- /dev/null +++ b/Pipfile @@ -0,0 +1,11 @@ +[[source]] +url = "https://pypi.org/simple" +verify_ssl = true +name = "pypi" + +[packages] + +[dev-packages] + +[requires] +python_version = "3.11" diff --git a/simplemind/client.py b/simplemind/client.py index 5bb1dbe..60e3211 100644 --- a/simplemind/client.py +++ b/simplemind/client.py @@ -3,6 +3,10 @@ from simplemind.models import Conversation, AIResponse from simplemind.concepts import Context from simplemind.integrations.openai import OpenAI from simplemind.integrations.anthropic import Anthropic +import logging + +logger = logging.getLogger(__name__) + class Client: def __init__(self, api_key: str, context: Optional[Context] = None): @@ -19,12 +23,24 @@ class Client: def create_conversation(self, provider: str = "openai") -> Conversation: if provider not in self.providers: raise ValueError(f"Provider '{provider}' not supported.") - return self.providers[provider].create_conversation(initial_message="Hello!", context=self.context.dict()) + return self.providers[provider].create_conversation( + initial_message="Hello!", context=self.context.dict() + ) - def send_message(self, conversation: Conversation, message: str, provider: str = "openai") -> AIResponse: + def _handle_api_error(self, error: Exception, operation: str): + """Handle API errors in a consistent way.""" + logger.error(f"Error during {operation}: {str(error)}") + raise RuntimeError(f"Failed to {operation}: {str(error)}") + + def send_message( + self, conversation: Conversation, message: str, provider: str = "openai" + ) -> AIResponse: if provider not in self.providers: raise ValueError(f"Provider '{provider}' not supported.") - return self.providers[provider].send_message(conversation.id, message) + try: + return self.providers[provider].send_message(conversation.id, message) + except Exception as e: + self._handle_api_error(e, "send message") @property def available_models(self): @@ -32,4 +48,3 @@ class Client: for name, provider in self.providers.items(): available[name] = provider.available_models return available - diff --git a/simplemind/concepts.py b/simplemind/concepts.py index 147d6a0..4e7da47 100644 --- a/simplemind/concepts.py +++ b/simplemind/concepts.py @@ -4,6 +4,7 @@ from simplemind.plugins.base import BasePlugin class Context(BaseModel): + model_config = {"arbitrary_types_allowed": True} plugins: Dict[str, BasePlugin] = {} def add_plugin(self, name: str, plugin: BasePlugin): diff --git a/simplemind/config.py b/simplemind/config.py index a14bb25..1a45fea 100644 --- a/simplemind/config.py +++ b/simplemind/config.py @@ -1,4 +1,4 @@ -from pydantic import BaseSettings +from pydantic_settings import BaseSettings class Settings(BaseSettings): diff --git a/simplemind/core.py b/simplemind/core.py index e69de29..8b11bb1 100644 --- a/simplemind/core.py +++ b/simplemind/core.py @@ -0,0 +1,44 @@ +from typing import Dict, Any, Optional +from .models import AIResponse +from .concepts import Context +from .integrations.base import BaseClientProvider + + +class SimpleMind: + """Main class for SimpleMind functionality.""" + + def __init__( + self, api_key: str, provider: str = "openai", context: Optional[Context] = None + ): + """Initialize SimpleMind with the specified provider.""" + self.api_key = api_key + self.provider = provider + self.context = context or Context() + self._client = self._get_provider() + + def _get_provider(self) -> BaseClientProvider: + """Get the appropriate provider client.""" + from .integrations.openai import OpenAI + from .integrations.anthropic import Anthropic + + providers = {"openai": OpenAI, "anthropic": Anthropic} + + if self.provider not in providers: + raise ValueError( + f"Provider '{self.provider}' not supported. Available providers: {list(providers.keys())}" + ) + + return providers[self.provider](api_key=self.api_key) + + def generate(self, prompt: str, **kwargs) -> AIResponse: + """Generate a response using the configured provider.""" + return self._client.message(prompt, **kwargs) + + def create_conversation(self, initial_message: str) -> str: + """Create a new conversation and return its ID.""" + conversation = self._client.create_conversation(initial_message) + return conversation.id + + def send_message(self, conversation_id: str, message: str) -> AIResponse: + """Send a message in an existing conversation.""" + return self._client.send_message(conversation_id, message) diff --git a/simplemind/integrations/anthropic.py b/simplemind/integrations/anthropic.py index 15678dc..ad532b0 100644 --- a/simplemind/integrations/anthropic.py +++ b/simplemind/integrations/anthropic.py @@ -9,6 +9,9 @@ from ..models import AIResponse, Conversation from ..logger import logger +DEFAULT_MODEL = "claude-3-5-sonnet-20240620" + + class Anthropic(BaseClientProvider): def __init__(self, model: str = DEFAULT_MODEL, api_key: Optional[str] = None): super().__init__(model=model, api_key=api_key) diff --git a/simplemind/integrations/base.py b/simplemind/integrations/base.py index 361adf9..40e8742 100644 --- a/simplemind/integrations/base.py +++ b/simplemind/integrations/base.py @@ -4,7 +4,6 @@ from pydantic import BaseModel from typing import Any, Dict, List, Optional from ..models import AIResponse, Conversation, Message import uuid -from abc import ABC, abstractmethod DEFAULT_MODEL = "gpt-4o" @@ -19,31 +18,27 @@ class BaseClientProvider: self._api_key = api_key self.conversations: Dict[str, Conversation] = {} - @abstractmethod def login(self): """Initializes the AI provider client.""" msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - @abstractmethod def test_connection(self) -> bool: """Tests the connection to the AI provider client.""" msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - @abstractmethod def health_check(self): """Checks the health of the AI provider client.""" msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - @abstractmethod + @property def available_models(self) -> List[str]: """Returns the available models from the AI provider client.""" msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - @abstractmethod def message(self, message: str, **kwargs) -> AIResponse: """Generates a response from the AI provider client.""" msg = "This method must be implemented by the AI provider client." @@ -95,7 +90,9 @@ class BaseClientProvider: conversation.context.update(context_update) response = self.generate_response(conversation) - conversation.messages.append(Message(role="assistant", content=response.text)) + conversation.messages.append( + Message(role="assistant", content=response.choices[0].message.content) + ) return response def generate_response(self, conversation: Conversation) -> AIResponse: @@ -108,9 +105,3 @@ class BaseClientProvider: if conversation_id not in self.conversations: raise ValueError("Conversation ID does not exist.") return self.conversations[conversation_id] - - -class BasePlugin(ABC): - @abstractmethod - def execute(self, context, *args, **kwargs): - pass diff --git a/simplemind/integrations/openai.py b/simplemind/integrations/openai.py index d67d9f5..cd808da 100644 --- a/simplemind/integrations/openai.py +++ b/simplemind/integrations/openai.py @@ -8,6 +8,10 @@ from ..models import AIResponse, Conversation from ..logger import logger from simplemind.config import settings + +DEFAULT_MODEL = "gpt-4o" + + class OpenAI(BaseClientProvider): def __init__(self, model: str = DEFAULT_MODEL, api_key: Optional[str] = None): super().__init__(model=model, api_key=api_key) @@ -46,22 +50,25 @@ class OpenAI(BaseClientProvider): return False def generate_response(self, conversation: Conversation) -> AIResponse: - messages = [ - {"role": msg.role, "content": msg.content} for msg in conversation.messages - ] + messages = conversation.get_messages() params = { - "messages": messages, "model": self.model, + "messages": [ + { + "role": msg.role, + "content": [{"type": "text", "text": msg.content}], # New format + } + for msg in messages + ], + "temperature": getattr( + self, "temperature", 0.7 + ), # Use 0.7 as default if not set } - if conversation.context: - params["context"] = conversation.context try: completion = self.client.chat.completions.create(**params) - response_text = completion.choices[0].message.content - metadata = {"model": completion.model, "usage": completion.usage} - logger.info("Generated response from OpenAI.") - return AIResponse(text=response_text, response=completion, metadata=metadata) + return completion except Exception as e: - logger.error(f"Error generating response: {e}") - raise e + # Enhanced error handling (optional) + logger.error(f"OpenAI API Error: {e}") + raise RuntimeError(f"Failed to generate response: {e}") diff --git a/simplemind/models.py b/simplemind/models.py index ea097e5..91fc603 100644 --- a/simplemind/models.py +++ b/simplemind/models.py @@ -1,6 +1,7 @@ from pydantic import BaseModel from typing import Any, Dict, List, Optional import uuid +from datetime import datetime class AIRequest(BaseModel): @@ -21,14 +22,27 @@ class AIResponse(BaseModel): class Message(BaseModel): - role: str # "user" or "assistant" + role: str # "user", "assistant", "system" content: str + created_at: datetime = datetime.now() class Conversation(BaseModel): id: str messages: List[Message] = [] - context: Optional[Dict[str, Any]] = {} + created_at: datetime = datetime.now() + updated_at: datetime = datetime.now() + + def get_messages(self) -> List[Message]: + """Returns a list of messages in the conversation.""" + return self.messages + + def add_message(self, role: str, content: str) -> Message: + """Adds a new message to the conversation.""" + message = Message(role=role, content=content) + self.messages.append(message) + self.updated_at = datetime.now() + return message class ConversationRequest(BaseModel): diff --git a/simplemind/plugins/base.py b/simplemind/plugins/base.py new file mode 100644 index 0000000..5aa48e5 --- /dev/null +++ b/simplemind/plugins/base.py @@ -0,0 +1,29 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict + + +class BasePlugin(ABC): + """Base class for all SimpleMind plugins.""" + + def __init__(self): + self.is_enabled = True + + @abstractmethod + def process(self, context: Dict[str, Any]) -> Dict[str, Any]: + """Process the context and return modified context. + + Args: + context: The current conversation context + + Returns: + Modified context dictionary + """ + pass + + def enable(self): + """Enable the plugin.""" + self.is_enabled = True + + def disable(self): + """Disable the plugin.""" + self.is_enabled = False diff --git a/simplemind/plugins/kv.py b/simplemind/plugins/kv.py index 136ca58..a83ae4b 100644 --- a/simplemind/plugins/kv.py +++ b/simplemind/plugins/kv.py @@ -1,10 +1,18 @@ -from .base import BasePlugin +from simplemind.plugins.base import BasePlugin class KVPlugin(BasePlugin): def __init__(self): self.store = {} - def execute(self, context, key, value): + def process(self, key: str, value=None): + """ + Get or set a value in the key-value store. + If value is None, returns the value for the key. + If value is provided, sets the value for the key and returns it. + """ + if value is None: + return self.store.get(key) + self.store[key] = value - return self.store + return value diff --git a/t.py b/t.py index ff34f6b..d5e0ca5 100644 --- a/t.py +++ b/t.py @@ -1,3 +1,4 @@ +import os from pprint import pprint from pydantic import BaseModel import simplemind @@ -8,16 +9,19 @@ from simplemind.chains.reverse_text import ReverseTextChain from simplemind.client import Client -class MyContext(Context): +class CustomContext(Context): def __init__(self): super().__init__() self.add_plugin("kv", KVPlugin()) - self.add_plugin("basic_memory", BasicMemoryPlugin()) + # self.add_plugin("basic_memory", BasicMemoryPlugin()) # Initialize context and client -context = MyContext() -aiclient = Client(api_key="YOUR_API_KEY", context=context) +ctx = CustomContext() +aiclient = Client( + context=ctx, + api_key=os.environ["OPENAI_API_KEY"], +) # Test connection and available models print(aiclient.available_models) From 2fc24bc94915692dd854e27d9fda9a22b24edfda Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 28 Oct 2024 09:12:38 -0400 Subject: [PATCH 3/5] Update simplemind/chains/base.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- simplemind/chains/base.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/simplemind/chains/base.py b/simplemind/chains/base.py index 279fdf5..41013c7 100644 --- a/simplemind/chains/base.py +++ b/simplemind/chains/base.py @@ -2,6 +2,23 @@ from abc import ABC, abstractmethod class BaseChain(ABC): + """Abstract base class for implementing chain operations. + + A Chain represents a processing step that can be executed on input data + and should be implemented by concrete classes to define specific behaviors. + """ + @abstractmethod - def run(self, input_data): + def run(self, input_data: str) -> str: + """Execute the chain's operation on the input data. + + Args: + input_data: The input string to be processed by the chain. + + Returns: + The processed output string. + + Raises: + ValueError: If the input data is invalid or cannot be processed. + """ pass From 60e52d15b3898170f87a4cea733cb3821d5d62cc Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 28 Oct 2024 09:12:51 -0400 Subject: [PATCH 4/5] Update simplemind/chains/reverse_text.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- simplemind/chains/reverse_text.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/simplemind/chains/reverse_text.py b/simplemind/chains/reverse_text.py index 7f102a2..cccdcbc 100644 --- a/simplemind/chains/reverse_text.py +++ b/simplemind/chains/reverse_text.py @@ -2,5 +2,24 @@ from .base import BaseChain class ReverseTextChain(BaseChain): - def run(self, input_data): + """Chain that reverses input text. + + This chain takes a text input and returns it reversed. For example, + "hello" becomes "olleh". + """ + + def run(self, input_data: str) -> str: + """Reverse the input text. + + Args: + input_data: The text to reverse. + + Returns: + The reversed text. + + Raises: + TypeError: If input_data is not a string. + """ + if not isinstance(input_data, str): + raise TypeError("Input must be a string") return input_data[::-1] From 75fd12e180c2dcd92f15f460c8d0f289825da574 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 28 Oct 2024 16:04:08 -0400 Subject: [PATCH 5/5] Delete Pipfile --- Pipfile | 11 ----------- 1 file changed, 11 deletions(-) delete mode 100644 Pipfile diff --git a/Pipfile b/Pipfile deleted file mode 100644 index 0757494..0000000 --- a/Pipfile +++ /dev/null @@ -1,11 +0,0 @@ -[[source]] -url = "https://pypi.org/simple" -verify_ssl = true -name = "pypi" - -[packages] - -[dev-packages] - -[requires] -python_version = "3.11"