diff --git a/autogpt/agent/agent.py b/autogpt/agent/agent.py index fd6a7cb6e1..9e0f54e36b 100644 --- a/autogpt/agent/agent.py +++ b/autogpt/agent/agent.py @@ -1,11 +1,10 @@ from colorama import Fore, Style from autogpt.app import execute_command, get_command -from autogpt.chat import chat_with_ai, create_chat_message from autogpt.config import Config from autogpt.json_utils.json_fix_llm import fix_json_using_multiple_techniques from autogpt.json_utils.utilities import validate_json -from autogpt.llm_utils import create_chat_completion +from autogpt.llm import chat_with_ai, create_chat_completion, create_chat_message from autogpt.logs import logger, print_assistant_thoughts from autogpt.speech import say_text from autogpt.spinner import Spinner diff --git a/autogpt/agent/agent_manager.py b/autogpt/agent/agent_manager.py index 257127a303..1283fdaefb 100644 --- a/autogpt/agent/agent_manager.py +++ b/autogpt/agent/agent_manager.py @@ -4,7 +4,7 @@ from __future__ import annotations from typing import List from autogpt.config.config import Config -from autogpt.llm_utils import create_chat_completion +from autogpt.llm import create_chat_completion from autogpt.singleton import Singleton from autogpt.types.openai import Message diff --git a/autogpt/commands/analyze_code.py b/autogpt/commands/analyze_code.py index 47cfc1e765..cff8a737e1 100644 --- a/autogpt/commands/analyze_code.py +++ b/autogpt/commands/analyze_code.py @@ -2,7 +2,7 @@ from __future__ import annotations from autogpt.commands.command import command -from autogpt.llm_utils import call_ai_function +from autogpt.llm import call_ai_function @command( diff --git a/autogpt/commands/improve_code.py b/autogpt/commands/improve_code.py index f953cf2bd3..ff4cf48c0d 100644 --- a/autogpt/commands/improve_code.py +++ b/autogpt/commands/improve_code.py @@ -3,7 +3,7 @@ from __future__ import annotations import json from autogpt.commands.command import command -from autogpt.llm_utils import call_ai_function +from autogpt.llm import call_ai_function @command( diff --git a/autogpt/commands/write_tests.py b/autogpt/commands/write_tests.py index 91cd930429..c6a0332474 100644 --- a/autogpt/commands/write_tests.py +++ b/autogpt/commands/write_tests.py @@ -4,7 +4,7 @@ from __future__ import annotations import json from autogpt.commands.command import command -from autogpt.llm_utils import call_ai_function +from autogpt.llm import call_ai_function @command( diff --git a/autogpt/json_utils/json_fix_llm.py b/autogpt/json_utils/json_fix_llm.py index 10317accb9..73f766480d 100644 --- a/autogpt/json_utils/json_fix_llm.py +++ b/autogpt/json_utils/json_fix_llm.py @@ -11,7 +11,7 @@ from regex import regex from autogpt.config import Config from autogpt.json_utils.json_fix_general import correct_json -from autogpt.llm_utils import call_ai_function +from autogpt.llm import call_ai_function from autogpt.logs import logger from autogpt.speech import say_text diff --git a/autogpt/llm/__init__.py b/autogpt/llm/__init__.py new file mode 100644 index 0000000000..3a958285de --- /dev/null +++ b/autogpt/llm/__init__.py @@ -0,0 +1,22 @@ +from autogpt.llm.api_manager import ApiManager +from autogpt.llm.chat import chat_with_ai, create_chat_message, generate_context +from autogpt.llm.llm_utils import ( + call_ai_function, + create_chat_completion, + get_ada_embedding, +) +from autogpt.llm.modelsinfo import COSTS +from autogpt.llm.token_counter import count_message_tokens, count_string_tokens + +__all__ = [ + "ApiManager", + "create_chat_message", + "generate_context", + "chat_with_ai", + "call_ai_function", + "create_chat_completion", + "get_ada_embedding", + "COSTS", + "count_message_tokens", + "count_string_tokens", +] diff --git a/autogpt/api_manager.py b/autogpt/llm/api_manager.py similarity index 98% rename from autogpt/api_manager.py rename to autogpt/llm/api_manager.py index b16f3af710..9143389e88 100644 --- a/autogpt/api_manager.py +++ b/autogpt/llm/api_manager.py @@ -3,8 +3,8 @@ from __future__ import annotations import openai from autogpt.config import Config +from autogpt.llm.modelsinfo import COSTS from autogpt.logs import logger -from autogpt.modelsinfo import COSTS from autogpt.singleton import Singleton diff --git a/autogpt/chat.py b/autogpt/llm/chat.py similarity index 94% rename from autogpt/chat.py rename to autogpt/llm/chat.py index 469ec9bd38..9bfd37e025 100644 --- a/autogpt/chat.py +++ b/autogpt/llm/chat.py @@ -3,10 +3,10 @@ from random import shuffle from openai.error import RateLimitError -from autogpt import token_counter -from autogpt.api_manager import ApiManager from autogpt.config import Config -from autogpt.llm_utils import create_chat_completion +from autogpt.llm.api_manager import ApiManager +from autogpt.llm.llm_utils import create_chat_completion +from autogpt.llm.token_counter import count_message_tokens from autogpt.logs import logger from autogpt.types.openai import Message @@ -43,7 +43,7 @@ def generate_context(prompt, relevant_memory, full_message_history, model): next_message_to_add_index = len(full_message_history) - 1 insertion_index = len(current_context) # Count the currently used tokens - current_tokens_used = token_counter.count_message_tokens(current_context, model) + current_tokens_used = count_message_tokens(current_context, model) return ( next_message_to_add_index, current_tokens_used, @@ -114,7 +114,7 @@ def chat_with_ai( prompt, relevant_memory, full_message_history, model ) - current_tokens_used += token_counter.count_message_tokens( + current_tokens_used += count_message_tokens( [create_chat_message("user", user_input)], model ) # Account for user input (appended later) @@ -122,9 +122,7 @@ def chat_with_ai( # print (f"CURRENT TOKENS USED: {current_tokens_used}") message_to_add = full_message_history[next_message_to_add_index] - tokens_to_add = token_counter.count_message_tokens( - [message_to_add], model - ) + tokens_to_add = count_message_tokens([message_to_add], model) if current_tokens_used + tokens_to_add > send_token_limit: break @@ -175,7 +173,7 @@ def chat_with_ai( ) if not plugin_response or plugin_response == "": continue - tokens_to_add = token_counter.count_message_tokens( + tokens_to_add = count_message_tokens( [create_chat_message("system", plugin_response)], model ) if current_tokens_used + tokens_to_add > send_token_limit: diff --git a/autogpt/llm_utils.py b/autogpt/llm/llm_utils.py similarity index 99% rename from autogpt/llm_utils.py rename to autogpt/llm/llm_utils.py index 11f8da3a16..c1ba5fa5e0 100644 --- a/autogpt/llm_utils.py +++ b/autogpt/llm/llm_utils.py @@ -8,8 +8,8 @@ import openai from colorama import Fore, Style from openai.error import APIError, RateLimitError, Timeout -from autogpt.api_manager import ApiManager from autogpt.config import Config +from autogpt.llm.api_manager import ApiManager from autogpt.logs import logger from autogpt.types.openai import Message diff --git a/autogpt/modelsinfo.py b/autogpt/llm/modelsinfo.py similarity index 100% rename from autogpt/modelsinfo.py rename to autogpt/llm/modelsinfo.py diff --git a/autogpt/token_counter.py b/autogpt/llm/token_counter.py similarity index 100% rename from autogpt/token_counter.py rename to autogpt/llm/token_counter.py diff --git a/autogpt/logs.py b/autogpt/logs.py index fc529f25b3..c1ac126cde 100644 --- a/autogpt/logs.py +++ b/autogpt/logs.py @@ -1,11 +1,9 @@ """Logging module for Auto-GPT.""" -import json import logging import os import random import re import time -import traceback from logging import LogRecord from colorama import Fore, Style diff --git a/autogpt/memory/local.py b/autogpt/memory/local.py index 87d83cc596..18b7255cb5 100644 --- a/autogpt/memory/local.py +++ b/autogpt/memory/local.py @@ -7,7 +7,7 @@ from typing import Any, List import numpy as np import orjson -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm import get_ada_embedding from autogpt.memory.base import MemoryProviderSingleton EMBED_DIM = 1536 diff --git a/autogpt/memory/milvus.py b/autogpt/memory/milvus.py index 41f9ed1309..40694df558 100644 --- a/autogpt/memory/milvus.py +++ b/autogpt/memory/milvus.py @@ -4,7 +4,7 @@ import re from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, connections from autogpt.config import Config -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm import get_ada_embedding from autogpt.memory.base import MemoryProviderSingleton diff --git a/autogpt/memory/pinecone.py b/autogpt/memory/pinecone.py index 3aa78f9edf..a0bd162ac5 100644 --- a/autogpt/memory/pinecone.py +++ b/autogpt/memory/pinecone.py @@ -1,7 +1,7 @@ import pinecone from colorama import Fore, Style -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm import get_ada_embedding from autogpt.logs import logger from autogpt.memory.base import MemoryProviderSingleton diff --git a/autogpt/memory/redismem.py b/autogpt/memory/redismem.py index 5618bb37f9..6b51277860 100644 --- a/autogpt/memory/redismem.py +++ b/autogpt/memory/redismem.py @@ -10,7 +10,7 @@ from redis.commands.search.field import TextField, VectorField from redis.commands.search.indexDefinition import IndexDefinition, IndexType from redis.commands.search.query import Query -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm import get_ada_embedding from autogpt.logs import logger from autogpt.memory.base import MemoryProviderSingleton diff --git a/autogpt/memory/weaviate.py b/autogpt/memory/weaviate.py index fbebbfd716..be29e1b98e 100644 --- a/autogpt/memory/weaviate.py +++ b/autogpt/memory/weaviate.py @@ -3,7 +3,7 @@ from weaviate import Client from weaviate.embedded import EmbeddedOptions from weaviate.util import generate_uuid5 -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm import get_ada_embedding from autogpt.memory.base import MemoryProviderSingleton diff --git a/autogpt/processing/text.py b/autogpt/processing/text.py index 9946951a7c..372e331d56 100644 --- a/autogpt/processing/text.py +++ b/autogpt/processing/text.py @@ -4,9 +4,8 @@ from typing import Dict, Generator, Optional import spacy from selenium.webdriver.remote.webdriver import WebDriver -from autogpt import token_counter from autogpt.config import Config -from autogpt.llm_utils import create_chat_completion +from autogpt.llm import count_message_tokens, create_chat_completion from autogpt.memory import get_memory CFG = Config() @@ -44,7 +43,7 @@ def split_text( ] expected_token_usage = ( - token_usage_of_chunk(messages=message_with_additional_sentence, model=model) + count_message_tokens(messages=message_with_additional_sentence, model=model) + 1 ) if expected_token_usage <= max_length: @@ -56,7 +55,7 @@ def split_text( create_message(" ".join(current_chunk), question) ] expected_token_usage = ( - token_usage_of_chunk(messages=message_this_sentence_only, model=model) + count_message_tokens(messages=message_this_sentence_only, model=model) + 1 ) if expected_token_usage > max_length: @@ -68,10 +67,6 @@ def split_text( yield " ".join(current_chunk) -def token_usage_of_chunk(messages, model): - return token_counter.count_message_tokens(messages, model) - - def summarize_text( url: str, text: str, question: str, driver: Optional[WebDriver] = None ) -> str: @@ -112,7 +107,7 @@ def summarize_text( memory.add(memory_to_add) messages = [create_message(chunk, question)] - tokens_for_chunk = token_counter.count_message_tokens(messages, model) + tokens_for_chunk = count_message_tokens(messages, model) print( f"Summarizing chunk {i + 1} / {len(chunks)} of length {len(chunk)} characters, or {tokens_for_chunk} tokens" ) diff --git a/autogpt/prompts/prompt.py b/autogpt/prompts/prompt.py index cbfcfca93d..7a53f60313 100644 --- a/autogpt/prompts/prompt.py +++ b/autogpt/prompts/prompt.py @@ -1,8 +1,8 @@ from colorama import Fore -from autogpt.api_manager import ApiManager from autogpt.config.ai_config import AIConfig from autogpt.config.config import Config +from autogpt.llm import ApiManager from autogpt.logs import logger from autogpt.prompts.generator import PromptGenerator from autogpt.setup import prompt_user diff --git a/autogpt/setup.py b/autogpt/setup.py index e3c251ecc2..ddc7c6bf43 100644 --- a/autogpt/setup.py +++ b/autogpt/setup.py @@ -6,7 +6,7 @@ from colorama import Fore, Style from autogpt import utils from autogpt.config import Config from autogpt.config.ai_config import AIConfig -from autogpt.llm_utils import create_chat_completion +from autogpt.llm import create_chat_completion from autogpt.logs import logger CFG = Config() diff --git a/tests/conftest.py b/tests/conftest.py index e3c590d288..da00058b7f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,8 +3,8 @@ from pathlib import Path import pytest from pytest_mock import MockerFixture -from autogpt.api_manager import ApiManager from autogpt.config import Config +from autogpt.llm import ApiManager from autogpt.workspace import Workspace pytest_plugins = ["tests.integration.agent_factory"] diff --git a/tests/integration/test_llm_utils.py b/tests/integration/test_llm_utils.py index 7f312bde19..74b860160f 100644 --- a/tests/integration/test_llm_utils.py +++ b/tests/integration/test_llm_utils.py @@ -3,7 +3,7 @@ import string import pytest from numpy.random import RandomState -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm.llm_utils import get_ada_embedding from tests.utils import requires_api_key diff --git a/tests/integration/weaviate_memory_tests.py b/tests/integration/weaviate_memory_tests.py index f408577554..2cea871c61 100644 --- a/tests/integration/weaviate_memory_tests.py +++ b/tests/integration/weaviate_memory_tests.py @@ -5,7 +5,7 @@ from weaviate import Client from weaviate.util import get_valid_uuid from autogpt.config import Config -from autogpt.llm_utils import get_ada_embedding +from autogpt.llm import get_ada_embedding from autogpt.memory.weaviate import WeaviateMemory diff --git a/tests/test_agent.py b/tests/test_agent.py index 9c67f7e41e..e859673159 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -3,10 +3,7 @@ from unittest.mock import MagicMock import pytest from autogpt.agent import Agent -from autogpt.chat import chat_with_ai from autogpt.config import Config -from autogpt.speech import say_text -from autogpt.utils import clean_input @pytest.fixture diff --git a/tests/test_agent_manager.py b/tests/test_agent_manager.py index d1bb3d8b73..e7e696ae5c 100644 --- a/tests/test_agent_manager.py +++ b/tests/test_agent_manager.py @@ -1,7 +1,7 @@ import pytest from autogpt.agent.agent_manager import AgentManager -from autogpt.llm_utils import create_chat_completion +from autogpt.llm import create_chat_completion @pytest.fixture diff --git a/tests/test_api_manager.py b/tests/test_api_manager.py index 71aa6d63b1..ba64a72f87 100644 --- a/tests/test_api_manager.py +++ b/tests/test_api_manager.py @@ -2,7 +2,7 @@ from unittest.mock import MagicMock, patch import pytest -from autogpt.api_manager import COSTS, ApiManager +from autogpt.llm import COSTS, ApiManager api_manager = ApiManager() diff --git a/tests/test_llm_utils.py b/tests/test_llm_utils.py index 83cac0de53..27c9ff883a 100644 --- a/tests/test_llm_utils.py +++ b/tests/test_llm_utils.py @@ -1,8 +1,8 @@ import pytest from openai.error import APIError, RateLimitError -from autogpt.llm_utils import get_ada_embedding, retry_openai_api -from autogpt.modelsinfo import COSTS +from autogpt.llm import COSTS, get_ada_embedding +from autogpt.llm.llm_utils import retry_openai_api @pytest.fixture(params=[RateLimitError, APIError]) @@ -19,7 +19,7 @@ def mock_create_embedding(mocker): mock_response.usage.prompt_tokens = 5 mock_response.__getitem__.side_effect = lambda key: [{"embedding": [0.1, 0.2, 0.3]}] return mocker.patch( - "autogpt.llm_utils.create_embedding", return_value=mock_response + "autogpt.llm.llm_utils.create_embedding", return_value=mock_response ) diff --git a/tests/test_token_counter.py b/tests/test_token_counter.py index ac25796daf..966a0aff11 100644 --- a/tests/test_token_counter.py +++ b/tests/test_token_counter.py @@ -1,7 +1,6 @@ import pytest -import tests.context -from autogpt.token_counter import count_message_tokens, count_string_tokens +from autogpt.llm import count_message_tokens, count_string_tokens def test_count_message_tokens(): diff --git a/tests/unit/test_chat.py b/tests/unit/test_chat.py index 61ea773c96..9cb7c625bd 100644 --- a/tests/unit/test_chat.py +++ b/tests/unit/test_chat.py @@ -2,7 +2,7 @@ import time from unittest.mock import patch -from autogpt.chat import create_chat_message, generate_context +from autogpt.llm import create_chat_message, generate_context def test_happy_path_role_content(): diff --git a/tests/unit/test_get_self_feedback.py b/tests/unit/test_get_self_feedback.py index 54641f9168..09707270d3 100644 --- a/tests/unit/test_get_self_feedback.py +++ b/tests/unit/test_get_self_feedback.py @@ -1,6 +1,6 @@ from autogpt.agent.agent import Agent from autogpt.config import AIConfig -from autogpt.llm_utils import create_chat_completion +from autogpt.llm import create_chat_completion def test_get_self_feedback(mocker):