From 2667583865e6c51e1a60c06f30683533f1c3d562 Mon Sep 17 00:00:00 2001 From: antonkulaga Date: Sun, 19 Jan 2025 21:03:08 +0200 Subject: [PATCH] improvements with tests and sqlite functions --- .github/workflows/run_tests.yaml | 4 + .gitignore | 3 +- .readthedocs.yaml | 19 ++- config/agent_profiles.yaml | 30 ---- docs/conf.py | 14 +- tests/test_chat_agent.py | 4 +- tests/test_pure_litellm.py | 4 +- tests/test_secretary_agent.py | 14 +- tests/test_stream.py | 2 + tools/just_agents/tools/db.py | 245 ++++++++++++++++++++++++++----- 10 files changed, 255 insertions(+), 84 deletions(-) delete mode 100644 config/agent_profiles.yaml diff --git a/.github/workflows/run_tests.yaml b/.github/workflows/run_tests.yaml index c868cdc..dd63779 100644 --- a/.github/workflows/run_tests.yaml +++ b/.github/workflows/run_tests.yaml @@ -16,6 +16,10 @@ jobs: env: GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }} + GROQ_API_KEY_1: ${{ secrets.GROQ_API_KEY_1 }} + GROQ_API_KEY_2: ${{ secrets.GROQ_API_KEY_2 }} + GROQ_API_KEY_3: ${{ secrets.GROQ_API_KEY_3 }} + GROQ_API_KEY_4: ${{ secrets.GROQ_API_KEY_4 }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} steps: diff --git a/.gitignore b/.gitignore index 3d4a1af..7585bf5 100644 --- a/.gitignore +++ b/.gitignore @@ -168,4 +168,5 @@ cython_debug/ *.egg-info -config/agent_profiles.yaml #to deal with testing bug \ No newline at end of file +./config/ #temporal config produced by tests +config/agent_profiles.yaml #to deal with testing bug diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 91a34df..e59cb6c 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -6,6 +6,8 @@ build: tools: python: "3.11" commands: + - pip install poetry + # First modify all pyproject.toml files to use exact versions - cp pyproject.toml pyproject.toml.bak - | sed -i \ @@ -16,12 +18,25 @@ build: -e "s|{ path = \"router\", develop = true }|\"0.4.9\"|g" \ -e "s|{ path = \"examples\", develop = true }|\"0.4.9\"|g" \ pyproject.toml + # Build packages in order following build.sh approach + # First build core + - cd core && poetry build && cd .. + - pip install core/dist/*.whl + # Build and install other packages + - cd tools && poetry build && cd .. && pip install tools/dist/*.whl + - cd coding && poetry build && cd .. && pip install coding/dist/*.whl + - cd web && poetry build && cd .. && pip install web/dist/*.whl + - cd router && poetry build && cd .. && pip install router/dist/*.whl + - cd examples && poetry build && cd .. && pip install examples/dist/*.whl + # Build the meta-package + - touch ./just_agents/__init__.py + - poetry build + - pip install dist/*.whl + - rm ./just_agents/__init__.py sphinx: configuration: docs/conf.py python: install: - - method: pip - path: . - requirements: docs/requirements.txt \ No newline at end of file diff --git a/config/agent_profiles.yaml b/config/agent_profiles.yaml deleted file mode 100644 index f4f7c23..0000000 --- a/config/agent_profiles.yaml +++ /dev/null @@ -1,30 +0,0 @@ -SecretaryAgent: - autoload_from_yaml: false - backstory: Developed to assist in understanding and documenting AI agents' capabilities - and characteristics. - class_qualname: just_agents.router.secretary_agent.SecretaryAgent - description: A skilled AI assistant focused on generating detailed profiles for - AI agents. - expertise_domain: AI agent analysis and description - extra_dict: - personality_traits: Agent's personality traits go here - goal: To provide accurate and informative profiles for various AI agents. - knowledge_sources: [] - limitations: Limited to analysis and description tasks; may not perform actions - outside of its defined scope. - llm_options: - model: gpt-4o-mini - temperature: 0.0 - model_name: gpt-4o-mini - personality_traits: Detail-oriented, analytical, concise, informative - role: AI assistant specializing in agent profiling. - system_prompt: |2- - - You are a skilled AI assistant specializing in analysis and description of AI agents. - You are tasked with generation of a minimalistic and concise yet detail-rich profile for an AI agent, based on the AVAILABLE_ATTRIBUTES, - including 'system_prompt', 'llm_options' and any other. Your task is to fill in values of a JSON-formatted profile - that matches the PROFILE_UPDATE_TEMPLATE provided below. Values of the template describe what output is expected for each field. - Only populate fields based on the well-established information, don't make up anything. - Double-check that the output contains only a valid JSON with all the fields specified in PROFILE_UPDATE_TEMPLATE. - Never include any additional text or explanations in your reply. - task: Generate concise and detail-rich profiles for AI agents. diff --git a/docs/conf.py b/docs/conf.py index f3991b2..d6ebb5f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -11,7 +11,8 @@ 'sphinx.ext.viewcode', 'myst_parser', 'sphinx.ext.intersphinx', - 'sphinx_autodoc_typehints' + 'sphinx_autodoc_typehints', + 'sphinx.ext.autoapi' ] # Add after the existing extensions @@ -20,8 +21,7 @@ 'pydantic', 'requests', 'numpydoc', - 'rich', - 'just_agents' + 'rich' ] # Theme settings @@ -52,4 +52,10 @@ myst_enable_extensions = [ "colon_fence", "deflist" -] \ No newline at end of file +] + +# AutoAPI settings +autoapi_type = 'python' +autoapi_dirs = ['../'] +autoapi_options = ['members', 'undoc-members', 'show-inheritance', 'show-module-summary'] +autoapi_python_use_implicit_namespaces = True \ No newline at end of file diff --git a/tests/test_chat_agent.py b/tests/test_chat_agent.py index 1705202..36dd0bc 100644 --- a/tests/test_chat_agent.py +++ b/tests/test_chat_agent.py @@ -53,7 +53,9 @@ def test_quering(open_genes_db): goal=f"help users by using SQL syntax to form comands to work with the {open_genes_db} sqlite database", task="formulate appropriate comands to operate in the given database.", tools=[sqlite_query], - llm_options=LLAMA3_3) + llm_options=LLAMA3_3, + key_list_env="GROQ_API_KEY" # LOAD GROQ_API_KEY FROM ENVIRONMENT VARIABLES + ) response = agent.query("Show me all tables in the database") diff --git a/tests/test_pure_litellm.py b/tests/test_pure_litellm.py index 41c0dc5..b537114 100644 --- a/tests/test_pure_litellm.py +++ b/tests/test_pure_litellm.py @@ -3,8 +3,6 @@ import pytest from dotenv import load_dotenv - - def get_current_weather(location: str): """ Gets the current weather in a given location @@ -58,7 +56,7 @@ def test_oai_works(load_env): OPENAI_GPT4oMINI,_ = load_env execute_completion(OPENAI_GPT4oMINI) -@pytest.mark.skip(reason="until fixed in https://github.com/BerriAI/litellm/issues/7621") +#@pytest.mark.skip(reason="until fixed in https://github.com/BerriAI/litellm/issues/7621") def test_grok_bug(load_env): _, LLAMA3_3 = load_env execute_completion(LLAMA3_3) \ No newline at end of file diff --git a/tests/test_secretary_agent.py b/tests/test_secretary_agent.py index 6ed882c..ccd03a3 100644 --- a/tests/test_secretary_agent.py +++ b/tests/test_secretary_agent.py @@ -4,16 +4,18 @@ import json import just_agents.llm_options from just_agents.router.secretary_agent import SecretaryAgent +from pytest import TempPathFactory, FixtureRequest +from typing import Tuple, Any @pytest.fixture(scope='module') -def temp_config_path(tmpdir_factory): +def temp_config_path(tmpdir_factory: TempPathFactory) -> str: # Create a temporary directory for YAML files dotenv.load_dotenv(override=True) tmpdir_factory.mktemp('config') return SecretaryAgent.DEFAULT_CONFIG_PATH @pytest.fixture -def secretary_autoload_false(temp_config_path): +def secretary_autoload_false(temp_config_path: str) -> Tuple[SecretaryAgent, bool]: dotenv.load_dotenv(override=True) params = { 'autoload_from_yaml': False, @@ -28,7 +30,7 @@ def secretary_autoload_false(temp_config_path): return secretary, result @pytest.fixture -def secretary_autoload_true(temp_config_path, secretary_autoload_false): +def secretary_autoload_true(temp_config_path: str, secretary_autoload_false: Tuple[SecretaryAgent, bool]) -> Tuple[SecretaryAgent, bool]: # Ensure the YAML file exists from the previous fixture dotenv.load_dotenv(override=True) params = { @@ -47,14 +49,14 @@ def secretary_autoload_true(temp_config_path, secretary_autoload_false): secretary.save_to_yaml() return secretary, result -def test_secretary_autoload_false(secretary_autoload_false): +def test_secretary_autoload_false(secretary_autoload_false: Tuple[SecretaryAgent, bool]) -> None: dotenv.load_dotenv(override=True) secretary, result = secretary_autoload_false assert result is True, "Failed to update profile when autoload_from_yaml is False." assert secretary.description != secretary.DEFAULT_DESCRIPTION, "Description was not updated." assert secretary.llm_model_name is not None, "LLM model name is None." -def test_secretary_autoload_true(secretary_autoload_true): +def test_secretary_autoload_true(secretary_autoload_true: Tuple[SecretaryAgent, bool]) -> None: dotenv.load_dotenv(override=True) secretary, result = secretary_autoload_true assert result is True, "Failed to update profile when autoload_from_yaml is True." @@ -62,7 +64,7 @@ def test_secretary_autoload_true(secretary_autoload_true): assert secretary.extras.get("personality_traits") is not None, "Personality traits were not set." assert secretary.llm_model_name is not None, "LLM model name is None." -def test_new_secretary(temp_config_path): +def test_new_secretary(temp_config_path: str) -> None: # Load the secretary from the YAML file created in previous tests dotenv.load_dotenv(override=True) new_secretary = SecretaryAgent.from_yaml( diff --git a/tests/test_stream.py b/tests/test_stream.py index 5880bda..7343271 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -5,6 +5,8 @@ from just_agents.base_agent import BaseAgent from just_agents.llm_options import LLMOptions, LLAMA3_3, LLAMA3_2_VISION, OPENAI_GPT4oMINI from just_agents.just_tool import JustToolsBus + + @pytest.fixture(scope="module", autouse=True) def load_env(): load_dotenv(override=True) diff --git a/tools/just_agents/tools/db.py b/tools/just_agents/tools/db.py index 1b6b490..d174010 100644 --- a/tools/just_agents/tools/db.py +++ b/tools/just_agents/tools/db.py @@ -1,50 +1,221 @@ -import sqlite3 -from eliot import start_action +from typing import Dict, List, Union, Optional +from eliot import start_action, Message from pathlib import Path -def sqlite_query(database: str, sql: str) -> str: +import sqlite3 + +def sqlite_query(database: Union[Path, str], sql: str) -> str: """Execute a SQL query and return results as formatted text. Args: - database (Path | str): Path to the SQLite database file + database (Union[Path, str]): Path to the SQLite database file sql (str): SQL query to execute Returns: str: Query results formatted as semicolon-separated text. First line contains column names, followed by rows of data. Returns empty string if no results found. + + Raises: + sqlite3.Error: If there's any database-related error """ - # Log the query execution using eliot - start_action(action_type="sqlite_query", database=database, sql=sql) - - # Convert string path to Path object if needed - if isinstance(database, str): - database = Path(database).absolute() - - # Establish database connection - conn = sqlite3.connect(database) - cursor = conn.cursor() - - cursor.execute(sql) - try: - # Fetch and format query results - rows = cursor.fetchall() - if rows is None or len(rows) == 0: - conn.close() - return "" - - # Extract column names from cursor description - names = [description[0] for description in cursor.description] - - # Format output with column names as first line - text = "; ".join(names) + "\n" + with start_action(action_type="sqlite_query", database=str(database), sql=sql) as action: + try: + # Convert string path to Path object if needed + db_path = Path(database).absolute() if isinstance(database, str) else database + + # Use context manager for database connection + with sqlite3.connect(db_path) as conn: + cursor = conn.cursor() + cursor.execute(sql) + + # Fetch and format query results + rows = cursor.fetchall() + if not rows: + action.log(message_type="query_result", status="empty") + return "" + + # Extract column names from cursor description + names = [description[0] for description in cursor.description] + + # Format output with column names as first line + text = "; ".join(names) + "\n" + + # Add data rows, converting all values to strings + text += "\n".join("; ".join(str(i) for i in row) for row in rows) + "\n" + + action.log(message_type="query_result", status="success", row_count=len(rows)) + return text + + except sqlite3.Error as e: + action.log(message_type="query_error", error=str(e)) + raise + +def sqlite_query_read_only(database: Union[Path, str], sql: str) -> str: + """Execute a read-only SQL query and return results as formatted text. + Prevents any database modifications. + + Args: + database (Union[Path, str]): Path to the SQLite database file + sql (str): SQL query to execute (SELECT queries only) + + Returns: + str: Query results formatted as semicolon-separated text. + First line contains column names, followed by rows of data. + Returns empty string if no results found. + + Raises: + ValueError: If the query contains potential modification statements + sqlite3.Error: If there's any database-related error + """ + with start_action(action_type="sqlite_query_read_only", database=str(database), sql=sql) as action: + # Check for potentially dangerous operations + sql_lower = sql.lower().strip() + forbidden_keywords = [ + 'insert', 'update', 'delete', 'drop', 'alter', 'create', + 'replace', 'truncate', 'attach', 'detach', 'vacuum' + ] - # Add data rows, converting all values to strings - for row in rows: - row = [str(i) for i in row] - text += "; ".join(row) + "\n" - finally: - # Ensure connection is closed even if an error occurs - conn.close() - - return text + if not sql_lower.startswith('select'): + raise ValueError("Only SELECT queries are allowed in read-only mode") + + for keyword in forbidden_keywords: + if keyword in sql_lower: + raise ValueError(f"Forbidden keyword '{keyword}' detected in query") + + try: + db_path = Path(database).absolute() if isinstance(database, str) else database + + # Use URI format to enable strict read-only mode + uri = f"file:{db_path}?mode=ro&immutable=1&nolock=1" + + with sqlite3.connect(uri, uri=True) as conn: + # Set additional safety restrictions + conn.execute("PRAGMA query_only = ON;") + conn.execute("PRAGMA temp_store = MEMORY;") + + cursor = conn.cursor() + cursor.execute(sql) + + # Rest of the function is same as sqlite_query + rows = cursor.fetchall() + if not rows: + action.log(message_type="query_result", status="empty") + return "" + + names = [description[0] for description in cursor.description] + text = "; ".join(names) + "\n" + text += "\n".join("; ".join(str(i) for i in row) for row in rows) + "\n" + + action.log(message_type="query_result", status="success", row_count=len(rows)) + return text + + except sqlite3.Error as e: + action.log(message_type="query_error", error=str(e)) + raise + +def extract_db_structure(db_path: Union[Path, str]) -> Dict[str, List[Dict]]: + """ + Extracts the structure (table names and column definitions) of all tables in the database. + + Args: + db_path (Union[Path, str]): The path to the SQLite database file. + + Returns: + Dict[str, List[Dict]]: A dictionary where keys are table names and values are lists of + dictionaries containing column information. + + Raises: + FileNotFoundError: If the database file does not exist. + ValueError: If the database file has an unsupported extension. + sqlite3.Error: If there's any database-related error + """ + db_path = Path(db_path) + + if not db_path.exists(): + raise FileNotFoundError(f"The database file '{db_path}' does not exist.") + if db_path.suffix not in {'.sqlite', '.db'}: + raise ValueError(f"Unsupported database file extension: '{db_path.suffix}'. Expected '.sqlite' or '.db'.") + + with start_action(action_type="extract_db_structure", database=str(db_path)) as action: + try: + with sqlite3.connect(db_path) as conn: + cursor = conn.cursor() + try: + # Query to get all table names in the database + cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") + tables = cursor.fetchall() + + # Loop through the tables and get their schema + db_structure = {} + for (table_name,) in tables: + # Get basic column info + cursor.execute(f"PRAGMA table_info('{table_name}');") + columns = cursor.fetchall() + + # Get foreign key information + cursor.execute(f"PRAGMA foreign_key_list('{table_name}');") + fk_data = cursor.fetchall() + + # Get index information + cursor.execute(f"PRAGMA index_list('{table_name}');") + indexes = cursor.fetchall() + + # Create a mapping of column name to its indexes + column_indexes = {} + for idx in indexes: + index_name = idx[1] + unique = bool(idx[2]) + cursor.execute(f"PRAGMA index_info('{index_name}');") + idx_info = cursor.fetchall() + for col_info in idx_info: + col_name = col_info[2] + column_indexes.setdefault(col_name, []).append({ + 'index_name': index_name, + 'unique': unique + }) + + # Create foreign key mapping + fk_mapping = {} + for fk in fk_data: + fk_column = fk[3] + fk_mapping[fk_column] = { + 'table': fk[2], + 'to_column': fk[4], + 'on_update': fk[5], + 'on_delete': fk[6] + } + + column_info = [] + for column in columns: + try: + col = { + 'cid': column[0], + 'name': column[1], + 'type': column[2], + 'notnull': bool(column[3]), + 'dflt_value': column[4], + 'pk': bool(column[5]), + 'foreign_key': fk_mapping.get(column[1]), + 'indexes': column_indexes.get(column[1], []), + 'hidden': False # SQLite 3.35.0+ supports hidden columns + } + column_info.append(col) + except IndexError as e: + action.log(message_type="column_parse_error", column=column, error=str(e)) + continue # Skip malformed column entries + + db_structure[table_name] = column_info + + action.log(message_type="structure_extracted", table_count=len(tables)) + return db_structure + + finally: + cursor.close() + + except sqlite3.Error as e: + action.log(message_type="structure_error", error=str(e)) + raise + except Exception as e: + action.log(message_type="unexpected_error", error=str(e)) + raise