From f9b1472cf40f61e009f5aed8b36c81d0329beef2 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Tue, 29 Oct 2024 06:37:03 -0400 Subject: [PATCH] examples dir --- examples/_context.py | 10 ++++ examples/contextual_memory.py | 99 +++++++++++++++++++++++++++++++++++ examples/generate_text.py | 5 ++ examples/requirements.txt | 4 ++ examples/simple_memory.py | 28 ++++++++++ simplemind/__init__.py | 21 +++----- 6 files changed, 154 insertions(+), 13 deletions(-) create mode 100644 examples/_context.py create mode 100644 examples/contextual_memory.py create mode 100644 examples/generate_text.py create mode 100644 examples/requirements.txt create mode 100644 examples/simple_memory.py diff --git a/examples/_context.py b/examples/_context.py new file mode 100644 index 0000000..c7f2c91 --- /dev/null +++ b/examples/_context.py @@ -0,0 +1,10 @@ +import os +import sys + +# Add the parent directory to the path so we can import the module. +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + + +import simplemind as sm + +__all__ = ["sm"] diff --git a/examples/contextual_memory.py b/examples/contextual_memory.py new file mode 100644 index 0000000..780fb40 --- /dev/null +++ b/examples/contextual_memory.py @@ -0,0 +1,99 @@ +from _context import sm + +from pydantic import BaseModel +import openai +import faiss +import numpy as np +import os +import pickle + + +class ContextualMemoryPlugin: + def __init__(self, api_key: str, memory_file: str = "memories.pkl", embedding_model: str = "text-embedding-ada-002"): + openai.api_key = api_key + self.memory_file = memory_file + self.embedding_model = embedding_model + self.memories = [] + self.embeddings = None + self.index = None + self.load_memories() + + def load_memories(self): + if os.path.exists(self.memory_file): + with open(self.memory_file, "rb") as f: + self.memories, self.embeddings = pickle.load(f) + self.build_faiss_index() + else: + self.memories = [] + self.embeddings = [] + self.index = faiss.IndexFlatL2(1536) # Dimension for ada-002 embeddings + + def save_memories(self): + with open(self.memory_file, "wb") as f: + pickle.dump((self.memories, self.embeddings), f) + + def build_faiss_index(self): + if self.embeddings: + self.index = faiss.IndexFlatL2(len(self.embeddings[0])) + self.index.add(np.array(self.embeddings).astype('float32')) + else: + self.index = faiss.IndexFlatL2(1536) + + def get_embedding(self, text: str) -> list: + response = openai.Embedding.create(input=text, model=self.embedding_model) + return response['data'][0]['embedding'] + + def add_memory(self, memory: str): + embedding = self.get_embedding(memory) + self.memories.append(memory) + self.embeddings.append(embedding) + self.index.add(np.array([embedding]).astype('float32')) + self.save_memories() + + def retrieve_memories(self, query: str, top_k: int = 3) -> list: + if not self.index or len(self.embeddings) == 0: + return [] + query_embedding = self.get_embedding(query) + D, I = self.index.search(np.array([query_embedding]).astype('float32'), top_k) + return [self.memories[i] for i in I[0] if i < len(self.memories)] + + def send_hook(self, conversation: sm.Conversation): + # Retrieve relevant memories based on the latest user message + if conversation.messages: + last_user_message = conversation.messages[-1].text + relevant_memories = self.retrieve_memories(last_user_message) + for memory in relevant_memories: + conversation.add_message(role="system", text=memory) + + def on_response(self, conversation: sm.Conversation, response: str): + # Optionally, add the AI's response to memories + self.add_memory(response) + +# Example Usage + +# Define a Pydantic model if needed +class Story(BaseModel): + title: str + content: str + +# Initialize the conversation with the ContextualMemoryPlugin +memory_plugin = ContextualMemoryPlugin(api_key=sm.settings.OPENAI_API_KEY) + +conversation = sm.create_conversation(llm_model="gpt-4o-mini", llm_provider="openai") +conversation.add_plugin(memory_plugin) + +# Add user message +conversation.add_message("user", "Tell me a story about a brave knight.") + +# Send the conversation and get the response +response = conversation.send() +print(response.text) + +# Optionally, retrieve structured data +structured_response = sm.generate_data( + "Summarize the above story.", + llm_model="gpt-4o", + llm_provider="openai", + response_model=Story, +) +print(structured_response) diff --git a/examples/generate_text.py b/examples/generate_text.py new file mode 100644 index 0000000..9be1020 --- /dev/null +++ b/examples/generate_text.py @@ -0,0 +1,5 @@ +from _context import sm + +r = sm.generate_text("Write a poem about the moon", llm_provider="openai", llm_model="gpt-3.5-turbo") + +print(r) diff --git a/examples/requirements.txt b/examples/requirements.txt new file mode 100644 index 0000000..ea40e41 --- /dev/null +++ b/examples/requirements.txt @@ -0,0 +1,4 @@ +numpy +openai +pydantic +faiss-cpu diff --git a/examples/simple_memory.py b/examples/simple_memory.py new file mode 100644 index 0000000..fbb79c1 --- /dev/null +++ b/examples/simple_memory.py @@ -0,0 +1,28 @@ +from _context import sm + + +class SimpleMemoryPlugin: + def __init__(self): + self.memories = [ + "the earth has fictionally beeen destroyed.", + "the moon is made of cheese.", + ] + + def yield_memories(self): + return (m for m in self.memories) + + def send_hook(self, conversation: sm.Conversation): + for m in self.yield_memories(): + conversation.add_message(role="system", text=m) + + +conversation = sm.create_conversation(llm_model="grok-beta", llm_provider="xai") +conversation.add_plugin(SimpleMemoryPlugin()) + +conversation.add_message( + role="user", + text="Write a poem about the moon", +) + +r = conversation.send() +print(r.text) diff --git a/simplemind/__init__.py b/simplemind/__init__.py index 2c7ca0e..f57bcd7 100644 --- a/simplemind/__init__.py +++ b/simplemind/__init__.py @@ -1,24 +1,17 @@ from .models import Conversation from .utils import find_provider - - -class SimpleMind: - - def structured_response( - self, prompt, *, llm_model=None, llm_provider=None, response_model=None - ): - provider = find_provider(llm_provider) - - return provider.structured_response( - llm_model=llm_model, response_model=response_model, prompt=prompt - ) +from .settings import settings def create_conversation(llm_model=None, llm_provider=None): + """Create a new conversation.""" + return Conversation(llm_model=llm_model, llm_provider=llm_provider) def generate_data(prompt, *, llm_model=None, llm_provider=None, response_model=None): + """Generate structured data from a given prompt.""" + provider = find_provider(llm_provider) return provider.structured_response( @@ -29,6 +22,8 @@ def generate_data(prompt, *, llm_model=None, llm_provider=None, response_model=N def generate_text(prompt, *, llm_model=None, llm_provider=None): + """Generate text from a given prompt.""" + provider = find_provider(llm_provider) return provider.generate_text(prompt=prompt, llm_model=llm_model) @@ -36,9 +31,9 @@ def generate_text(prompt, *, llm_model=None, llm_provider=None): __all__ = [ "Conversation", - "SimpleMind", "create_conversation", "find_provider", "generate_data", "generate_text", + "settings" ]