Merge pull request #1 from kennethreitz/refactor

WIP
This commit is contained in:
2024-10-28 16:05:16 -04:00
committed by GitHub
25 changed files with 588 additions and 174 deletions
+13
View File
@@ -109,3 +109,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:**
View File
+7
View File
@@ -0,0 +1,7 @@
from abc import ABC, abstractmethod
class BaseAgent(ABC):
@abstractmethod
def decide(self, context, *args, **kwargs):
pass
View File
+24
View File
@@ -0,0 +1,24 @@
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: 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
+25
View File
@@ -0,0 +1,25 @@
from .base import BaseChain
class ReverseTextChain(BaseChain):
"""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]
+50
View File
@@ -0,0 +1,50 @@
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
import logging
logger = logging.getLogger(__name__)
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 _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.")
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):
available = {}
for name, provider in self.providers.items():
available[name] = provider.available_models
return available
+15 -4
View File
@@ -1,6 +1,17 @@
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):
model_config = {"arbitrary_types_allowed": True}
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.")
+13
View File
@@ -0,0 +1,13 @@
from pydantic_settings 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()
+40 -17
View File
@@ -1,21 +1,44 @@
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.")
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."""
@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))
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()
@app.get("/health")
def health_check():
return {"status": "healthy"}
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)
+52 -21
View File
@@ -1,41 +1,72 @@
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
DEFAULT_MODEL = "claude-3-5-sonnet-20240620"
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
+53 -19
View File
@@ -1,6 +1,9 @@
# import logging
from pydantic import BaseModel
from typing import Any, Dict, List, Optional
from ..models import AIResponse, Conversation, Message
import uuid
DEFAULT_MODEL = "gpt-4o"
@@ -8,55 +11,42 @@ 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] = {}
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):
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)
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):
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):
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 +61,47 @@ 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.choices[0].message.content)
)
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]
+53 -47
View File
@@ -1,68 +1,74 @@
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
DEFAULT_MODEL = "gpt-4o"
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 = conversation.get_messages()
params = {
"messages": [{"role": "user", "content": message}],
"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
}
params.update(kwargs)
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)
return completion
except Exception as e:
# Enhanced error handling (optional)
logger.error(f"OpenAI API Error: {e}")
raise RuntimeError(f"Failed to generate response: {e}")
+15
View File
@@ -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")
+41 -3
View File
@@ -1,10 +1,12 @@
from pydantic import BaseModel
from typing import Any, ClassVar
from typing import Any, Dict, List, Optional
import uuid
from datetime import datetime
class AIRequest(BaseModel):
text: str
parameters: dict = {}
parameters: Dict[str, Any] = {}
def __str__(self):
return self.text
@@ -13,7 +15,43 @@ 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", "assistant", "system"
content: str
created_at: datetime = datetime.now()
class Conversation(BaseModel):
id: str
messages: List[Message] = []
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):
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] = {}
View File
+29
View File
@@ -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
+10
View File
@@ -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
+18
View File
@@ -0,0 +1,18 @@
from simplemind.plugins.base import BasePlugin
class KVPlugin(BasePlugin):
def __init__(self):
self.store = {}
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 value
View File
+19
View File
@@ -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
+28 -44
View File
@@ -1,54 +1,38 @@
import os
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 CustomContext(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
ctx = CustomContext()
aiclient = Client(
context=ctx,
api_key=os.environ["OPENAI_API_KEY"],
)
# 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
+54 -18
View File
@@ -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)
+3 -1
View File
@@ -1,3 +1,4 @@
import simplemind
aiclient = simplemind.Ollama()
@@ -25,4 +26,5 @@ conversation.say("What number did I ask you to remember?")
# Get the AI's response
reply = conversation.get_reply()
print(reply)
print(reply)
+26
View File
@@ -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()