diff --git a/bondai/models/openai/openai_embedding_model.py b/bondai/models/openai/openai_embedding_model.py index 85cf17c..64f2e5c 100644 --- a/bondai/models/openai/openai_embedding_model.py +++ b/bondai/models/openai/openai_embedding_model.py @@ -35,7 +35,7 @@ def max_tokens(self) -> int: def create_embedding(self, prompt: str) -> List[float] | List[List[float]]: return create_embedding( - prompt, self._model, connection_params=self._connection_params + prompt, connection_params=self._connection_params, model=self._model ) def count_tokens(self, prompt: str) -> int: diff --git a/bondai/models/openai/openai_wrapper.py b/bondai/models/openai/openai_wrapper.py index 2543b8e..08f9c67 100644 --- a/bondai/models/openai/openai_wrapper.py +++ b/bondai/models/openai/openai_wrapper.py @@ -69,8 +69,8 @@ def calculate_cost(model_name: str, usage: Dict): else: embedding_tokens += token_count embedding_costs += token_count * model["price_per_token"] - else: - print(f"Unknown model: {model_name}") + # else: + # print(f"Unknown model: {model_name}") def get_max_tokens(model: str) -> int: diff --git a/tests/memory/single_agent_with_memory.py b/tests/memory/single_agent_with_memory.py index c6eea34..46dbb8c 100644 --- a/tests/memory/single_agent_with_memory.py +++ b/tests/memory/single_agent_with_memory.py @@ -1,35 +1,61 @@ -from datetime import datetime -from bondai.models.openai import get_total_cost -from bondai.tools.file import FileWriteTool -from bondai.agents import ConversationalAgent -from bondai.memory import MemoryManager -from util import extract_text_from_directory - -memory_manager = MemoryManager() +from bondai.agents import ConversationalAgent, AgentEventNames +from bondai.models.openai import OpenAIEmbeddingModel +from bondai.memory import ( + MemoryManager, + InMemoryCoreMemoryDataSource, + InMemoryArchivalMemoryDataSource, +) + +import io +import requests +from PyPDF2 import PdfReader +from bondai.util import split_text + + +def retrieve_and_parse_pdf(url): + response = requests.get(url) + if response.status_code == 200: + pdf = PdfReader(io.BytesIO(response.content)) + text = "" + for page in pdf.pages: + text += page.extract_text() + "\n" + + return split_text(OpenAIEmbeddingModel(), text) + else: + return f"Error retrieving PDF: {response.status_code}" + + +memory_manager = MemoryManager( + core_memory_datasource=InMemoryCoreMemoryDataSource(), + archival_memory_datasource=InMemoryArchivalMemoryDataSource(), +) + memory_manager.core_memory.set( "user", "Name is George. Lives in New York. Has a dog named Max." ) memory_manager.archival_memory.insert_bulk( - extract_text_from_directory("./tests/memory/documents") + retrieve_and_parse_pdf("https://arxiv.org/pdf/2310.10501.pdf") ) -agent = ConversationalAgent(tools=[FileWriteTool()], memory_manager=memory_manager) - -message = "Start the conversation by sending the first message. You can exit any time by typing 'exit'." -while True: - user_input = input(message + "\n") - if user_input.lower() == "exit": - break - response = agent.send_message(user_input) - response.success = True - response.completed_at = datetime.now() - - if response: - message = response.message - if message.lower() == "exit": - break - else: - print("The agent has exited the conversation.") - break +agent = ConversationalAgent(memory_manager=memory_manager) +agent.on( + AgentEventNames.TOOL_COMPLETED, + lambda _, m: print( + f"*************\nTool: {m.tool_name}({str(m.tool_arguments)})\nOutput: {m.tool_output}\n\n" + ), +) -print(f"Total Cost: {get_total_cost()}") +response = agent.send_message("Do you know my name?") +print(response.message) + +response = agent.send_message("Actually my name is Kevin.") +print(response.message) + +response = agent.send_message( + ( + "Can you check your archival memory to see what information you have about Nemo Guardrails? " + "I'd like a full summary of the information you have about the project including an example " + "that demonstrates how to use Colang." + ) +) +print(response.message) diff --git a/tests/vision/single_agent_with_vision.py b/tests/vision/single_agent_with_vision.py index 4cd2e2b..200e118 100644 --- a/tests/vision/single_agent_with_vision.py +++ b/tests/vision/single_agent_with_vision.py @@ -1,25 +1,8 @@ -from datetime import datetime -from bondai.models.openai import get_total_cost from bondai.tools.vision import ImageAnalysisTool -from bondai.agents import ConversationalAgent +from bondai.agents import Agent -agent = ConversationalAgent(tools=[ImageAnalysisTool()]) - -message = "Start the conversation by sending the first message. You can exit any time by typing 'exit'." -while True: - user_input = input(message + "\n") - if user_input.lower() == "exit": - break - response = agent.send_message(user_input) - response.success = True - response.completed_at = datetime.now() - - if response: - message = response.message - if message.lower() == "exit": - break - else: - print("The agent has exited the conversation.") - break - -print(f"Total Cost: {get_total_cost()}") +agent = Agent(tools=[ImageAnalysisTool()]) +result = agent.run( + "What kind of animal this is? https://www.forbes.com/advisor/wp-content/uploads/2023/09/getty_creative.jpeg-900x510.jpg" +) +result.tool_arguments["results"] diff --git a/website/docs/examples/online-research/online-research.md b/website/docs/examples/online-research/online-research.md index 43fc18d..c8ee4c5 100644 --- a/website/docs/examples/online-research/online-research.md +++ b/website/docs/examples/online-research/online-research.md @@ -22,7 +22,7 @@ Next we will create an Agent with 3 tools: - **FileWriteTool** Allows the Agent to save the results to a file. ```python -from bondai import Agent +from bondai.agents import Agent from bondai.tools.search import GoogleSearchTool from bondai.tools.website import WebsiteQueryTool from bondai.tools.file import FileWriteTool