wip(agents-api): Add session sql queries

[creatorrr]

User description

Signed-off-by: Diwank Singh Tomer [email protected]

---

PR Type

Enhancement, Tests

---

Description

• Migrated database operations from CozoDB to PostgreSQL for sessions and entries
• Added new session model fields: system_template and forward_tool_calls
• Implemented comprehensive PostgreSQL queries for session and entry CRUD operations
• Added query utilities with timeout support and improved error handling
• Enhanced database schema with new fields and triggers
• Updated test suite with async/await patterns and improved coverage
• Added explicit API exports and documentation improvements

---

Changes walkthrough 📝

	Relevant files
Enhancement	20 files Sessions.py Add system template and tool call forwarding to session models agents-api/agents_api/autogen/Sessions.py Added system_template field to session-related models Added forward_tool_calls boolean field to control tool call forwarding Updated models to include new fields in request/response schemas +40/-0    utils.py Add PostgreSQL query utilities and type definitions            agents-api/agents_api/queries/utils.py Added PostgreSQL query preparation utilities Added type definitions for SQL query arguments Enhanced error handling for database operations Added timeout support for queries +118/-46 create_entries.py Add PostgreSQL entry creation queries                                        agents-api/agents_api/queries/entries/create_entries.py Implemented entry creation with PostgreSQL queries Added support for bulk entry creation Added validation and error handling +193/-0  delete_entries.py Add PostgreSQL entry deletion queries                                        agents-api/agents_api/queries/entries/delete_entries.py Implemented entry deletion with PostgreSQL queries Added support for bulk entry deletion Added cascade deletion for related records +150/-0  create_or_update_session.py Add PostgreSQL session upsert queries                                        agents-api/agents_api/queries/sessions/create_or_update_session.py Implemented session upsert with PostgreSQL queries Added participant lookup handling Added validation for required fields +154/-0  create_session.py Add PostgreSQL session creation queries                                    agents-api/agents_api/queries/sessions/create_session.py Implemented session creation with PostgreSQL queries Added participant lookup creation Added validation for required fields +143/-0  Sessions.py Update session models in integrations service                        integrations-service/integrations/autogen/Sessions.py Added system_template field to session models Added forward_tool_calls boolean field Synchronized model changes with agents-api +40/-0    list_sessions.py Add PostgreSQL session listing queries                                      agents-api/agents_api/queries/sessions/list_sessions.py Implemented session listing with PostgreSQL queries Added filtering and pagination support Added sorting options +106/-0  list_entries.py Add PostgreSQL entry listing queries                                          agents-api/agents_api/queries/entries/list_entries.py Implemented entry listing with PostgreSQL queries Added filtering and pagination support Added source filtering options +121/-0  get_session.py Add PostgreSQL session retrieval queries                                  agents-api/agents_api/queries/sessions/get_session.py Implemented session retrieval with PostgreSQL queries Added participant lookup joins Added error handling for missing sessions +83/-0    patch_session.py Add PostgreSQL session patch queries                                          agents-api/agents_api/queries/sessions/patch_session.py Implemented session patching with PostgreSQL queries Added partial update support Added validation for update fields +89/-0    delete_session.py Add PostgreSQL session deletion queries                                    agents-api/agents_api/queries/sessions/delete_session.py Implemented session deletion with PostgreSQL queries Added cascade deletion for lookups Added transaction support +69/-0    count_sessions.py Add PostgreSQL session counting queries                                    agents-api/agents_api/queries/sessions/count_sessions.py Implemented session counting with PostgreSQL queries Added developer filtering Added error handling +55/-0    __init__.py Initialize session queries module                                                agents-api/agents_api/queries/sessions/init.py Added module initialization for session queries Exported all session query functions Added module documentation +30/-0    app.py Improve database connection pool handling                                agents-api/agents_api/app.py Added null checks for postgres pool creation and closing Removed unused imports (json, asyncpg) +6/-4      __init__.py Create new entries query module                                                    agents-api/agents_api/queries/entries/init.py Created new entries module for SQL query functions Added imports for entry-related operations Added module documentation +21/-0    Entries.py Add model field to BaseEntry class                                              agents-api/agents_api/autogen/Entries.py Added model field to BaseEntry with default value "gpt-4o-mini" +1/-0      000015_entries.up.sql Enhance entries table schema and triggers                                memory-store/migrations/000015_entries.up.sql Added 'developer' to chat_role enum Added tokenizer field and modified timestamp type Added trigger to update session's updated_at timestamp Fixed token count calculation variable naming +24/-7    000016_entry_relations.up.sql Improve entry relations leaf node handling                              memory-store/migrations/000016_entry_relations.up.sql Replaced leaf node enforcement with automatic leaf status updates Added new trigger function for managing leaf status +19/-15  models.tsp Add system template and tool call forwarding to Session    typespec/sessions/models.tsp Added system_template field to Session model Added forward_tool_calls boolean field +6/-0
Tests	3 files test_session_queries.py Migrate session tests from CozoDB to PostgreSQL                    tests/test_session_queries.py Rewrote session query tests to use PostgreSQL instead of CozoDB Added async/await pattern for database operations Added new tests for session CRUD operations with SQL queries Improved test coverage for error cases and edge scenarios test_entry_queries.py Migrate entry tests from CozoDB to PostgreSQL                        tests/test_entry_queries.py Migrated entry query tests from CozoDB to PostgreSQL Added async/await pattern for database operations Updated test assertions for SQL query responses Added error handling test cases fixtures.py Update test fixtures for PostgreSQL database                          agents-api/tests/fixtures.py Updated test fixtures to use PostgreSQL instead of CozoDB Added async/await pattern for database operations Simplified fixture cleanup and resource management +64/-64
Documentation	2 files __init__.py Add explicit exports for developer queries                              agents-api/agents_api/queries/developers/init.py Added all list to explicitly define public API +7/-0      __init__.py Add explicit exports for agent queries                                      agents-api/agents_api/queries/agents/init.py Added __all__ list to explicitly define public API for agent queries +10/-0
Configuration changes	1 files env.py Add query timeout environment configuration                            agents-api/agents_api/env.py Added query_timeout configuration with default value of 90.0 +2/-0

---

💡 PR-Agent usage: Comment /help "your question" on any pull request to receive relevant information

[qodo-merge-pro-for-open-source]

CI Failure Feedback 🧐

(Checks updated until commit 619f973)

Action: Typecheck

Failed stage: Typecheck [❌]

Failed test name: agents_api.common.utils.cozo

Failure summary: The pytype check failed due to multiple issues: 1. Missing dependency: Cannot find module pycozo which is required in agents_api/common/utils/cozo.py 2. Type checking errors: - Invalid type error in pytype matcher - Stray type comments found in test_workflow_routes.py (lines 65 and 114) These errors prevented the build from completing successfully.

Relevant error logs: 1: ##[group]Operating System 2: Ubuntu ... 1186: [18/355] check agents_api.autogen.Executions 1187: [19/355] check agents_api.autogen.Entries 1188: [20/355] check agents_api.autogen.Agents 1189: [21/355] check agents_api.clients.__init__ 1190: [22/355] check agents_api.activities.sync_items_remote 1191: [23/355] check agents_api.common.protocol.agents 1192: [24/355] check agents_api.common.utils.datetime 1193: [25/355] check agents_api.common.utils.cozo 1194: FAILED: /home/runner/work/julep/julep/agents-api/.pytype/pyi/agents_api/common/utils/cozo.pyi 1195: /home/runner/work/julep/julep/agents-api/.venv/bin/python -m pytype.main --disable pyi-error --imports_info /home/runner/work/julep/julep/agents-api/.pytype/imports/agents_api.common.utils.cozo.imports --module-name agents_api.common.utils.cozo --platform linux -V 3.12 -o /home/runner/work/julep/julep/agents-api/.pytype/pyi/agents_api/common/utils/cozo.pyi --analyze-annotated --nofail --none-is-not-bool --quick --strict-none-binding /home/runner/work/julep/julep/agents-api/agents_api/common/utils/cozo.py 1196: /home/runner/work/julep/julep/agents-api/agents_api/common/utils/cozo.py:9:1: error: in <module>: Can't find module 'pycozo'. [import-error] 1197: from pycozo import Client 1198: ~~~~~~~~~~~~~~~~~~~~~~~~~ 1199: For more details, see https://google.github.io/pytype/errors.html#import-error ... 1209: [35/355] check agents_api.common.exceptions.__init__ 1210: [36/355] check agents_api.common.storage_handler 1211: [37/355] check agents_api.dependencies.exceptions 1212: [38/355] check agents_api.clients.integrations 1213: [39/355] check agents_api.clients.litellm 1214: [40/355] check agents_api.routers.docs.router 1215: [41/355] check agents_api.routers.sessions.router 1216: [42/355] check agents_api.queries.utils 1217: ERROR:pytype.matcher Invalid type: <class 'pytype.abstract.function.ParamSpecMatch'> ... 1284: [109/355] check tests.test_task_queries 1285: [110/355] check tests.test_messages_truncation 1286: [111/355] check agents_api.queries.users.delete_user 1287: [112/355] check agents_api.autogen.__init__ 1288: [113/355] check agents_api.queries.users.create_or_update_user 1289: [114/355] check tests.sample_tasks.test_find_selector 1290: [115/355] check tests.test_sessions 1291: [116/355] check tests.test_workflow_routes 1292: /home/runner/work/julep/julep/agents-api/tests/test_workflow_routes.py:65:1: error: : Stray type comment: object [ignored-type-comment] 1293: # type: object~~~~~~~~~~~~~~~ 1294: # type: object 1295: /home/runner/work/julep/julep/agents-api/tests/test_workflow_routes.py:114:1: error: : Stray type comment: object [ignored-type-comment] 1296: # type: object~~~~~~~~~~~~~~~ 1297: # type: object 1298: For more details, see https://google.github.io/pytype/errors.html#ignored-type-comment ... 1332: [150/355] check agents_api.common.exceptions.users 1333: [151/355] check agents_api.dependencies.__init__ 1334: [152/355] check agents_api.queries.users.__init__ 1335: [153/355] check tests.test_user_routes 1336: [154/355] check tests.test_docs_queries 1337: [155/355] check agents_api.dependencies.auth 1338: [156/355] check agents_api.queries.sessions.__init__ 1339: [157/355] check agents_api.common.utils.debug 1340: ninja: build stopped: cannot make progress due to previous errors. 1341: Computing dependencies 1342: Generated API key since not set in the environment: 37371723677871526212764655913813 1343: Sentry DSN not found. Sentry will not be enabled. 1344: Analyzing 327 sources with 0 local dependencies 1345: Leaving directory '.pytype' 1346: ##[error]Process completed with exit code 1.

---

✨ CI feedback usage guide:

---

The CI feedback tool (/checks) automatically triggers when a PR has a failed check.
The tool analyzes the failed checks and provides several feedbacks:

• Failed stage
• Failed test name
• Failure summary
• Relevant error logs

In addition to being automatically triggered, the tool can also be invoked manually by commenting on a PR:

/checks "https://github.com/{repo_name}/actions/runs/{run_number}/job/{job_number}"

where {repo_name} is the name of the repository, {run_number} is the run number of the failed check, and {job_number} is the job number of the failed check.

Configuration options

• enable_auto_checks_feedback - if set to true, the tool will automatically provide feedback when a check is failed. Default is true.
• excluded_checks_list - a list of checks to exclude from the feedback, for example: ["check1", "check2"]. Default is an empty list.
• enable_help_text - if set to true, the tool will provide a help message with the feedback. Default is true.
• persistent_comment - if set to true, the tool will overwrite a previous checks comment with the new feedback. Default is true.
• final_update_message - if persistent_comment is true and updating a previous checks message, the tool will also create a new message: "Persistent checks updated to latest commit". Default is true.

See more information about the checks tool in the docs.

[gitguardian]

⚠️ GitGuardian has uncovered 1 secret following the scan of your pull request.

Please consider investigating the findings and remediating the incidents. Failure to do so may lead to compromising the associated services or software components.

🔎 Detected hardcoded secret in your pull request

GitGuardian id	GitGuardian status	Secret	Commit	Filename
14144715	Triggered	Generic Password	c85c2fe	memory-store/docker-compose.yml	View secret

🛠 Guidelines to remediate hardcoded secrets

1. Understand the implications of revoking this secret by investigating where it is used in your code.
2. Replace and store your secret safely. Learn here the best practices.
3. Revoke and rotate this secret.
4. If possible, rewrite git history. Rewriting git history is not a trivial act. You might completely break other contributing developers' workflow and you risk accidentally deleting legitimate data.

To avoid such incidents in the future consider

• following these best practices for managing and storing secrets including API keys and other credentials
• install secret detection on pre-commit to catch secret before it leaves your machine and ease remediation.

---

^{🦉 GitGuardian detects secrets in your source code to help developers and security teams secure the modern development process. You are seeing this because you or someone else with access to this repository has authorized GitGuardian to scan your pull request.}

[qodo-merge-pro-for-open-source]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

🔒 Security concerns SQL injection: The list_entries_query in agents-api/agents_api/queries/entries/list_entries.py uses string format() to interpolate sort_by and direction parameters directly into the SQL query. While these parameters are validated as Literals, this approach is still risky and could potentially be exploited. Consider using parameterized queries or a safer SQL query builder.

⚡ Recommended focus areas for review SQL Injection Risk The SQL query string interpolation in list_entries_query uses format() with sort_by and direction parameters. Even though these are validated as Literals, direct string interpolation should be avoided. Error Handling The prepare_pg_query_args function has a broad catch-all case that raises ValueError. Consider adding more specific error cases and messages. Performance Concern The session listing query performs a LEFT JOIN with array_agg which could be slow for large datasets. Consider pagination or limiting the aggregated arrays.

[ellipsis-dev]

👍 Looks good to me! Reviewed everything up to 619f973 in 1 minute and 34 seconds

More details

• Looked at 4858 lines of code in 58 files
• Skipped 0 files when reviewing.
• Skipped posting 8 drafted comments based on config settings.

1. agents-api/agents_api/app.py:10

• Draft comment:
  Consider using getattr to safely access app.state.postgres_pool to avoid potential AttributeError if the attribute is not initialized.

if not getattr(app.state, 'postgres_pool', None):

• Reason this comment was not posted:
  Decided after close inspection that this draft comment was likely wrong and/or not actionable:
  FastAPI's state object is designed for storing arbitrary state and is initialized by default. The .state attribute always exists and is safe to access. The current code's null check is checking if the pool exists, not if the attribute exists. Using getattr() wouldn't provide additional safety since AttributeError isn't actually a risk here.
  Maybe I'm wrong about FastAPI's state behavior. Maybe there are edge cases where state isn't initialized.
  FastAPI's documentation and source code confirm that state is always initialized as a simple object. The current code's intention is clearly to check for the pool's existence, not handle attribute errors.
  The comment should be deleted as it suggests a change that doesn't provide real value and misunderstands the purpose of the null check.

2. agents-api/agents_api/queries/agents/create_agent.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The create_agent.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

3. agents-api/agents_api/queries/agents/create_or_update_agent.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The create_or_update_agent.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

4. agents-api/agents_api/queries/agents/delete_agent.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The delete_agent.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

5. agents-api/agents_api/queries/agents/get_agent.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The get_agent.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

6. agents-api/agents_api/queries/agents/list_agents.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The list_agents.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

7. agents-api/agents_api/queries/agents/patch_agent.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The patch_agent.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

8. agents-api/agents_api/queries/agents/update_agent.py:15

• Draft comment:
  The variable agent_query is defined but not used. Consider removing it to clean up the code.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The update_agent.py file has a variable agent_query that is defined but not used. This might be a leftover from refactoring and should be removed to clean up the code.

Workflow ID: wflow_kEmSVhruL4I4QDhV

---

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

[qodo-merge-pro-for-open-source]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Score
Possible issue	Fix SQL syntax error by removing trailing comma in column list The entry_relation_query has an extra comma after 'tail' which will cause a SQL syntax error. Remove the trailing comma in the column list. agents-api/agents_api/queries/entries/create_entries.py [46-54] entry_relation_query = """ INSERT INTO entry_relations ( session_id, head, relation, - tail, -) VALUES ($1, $2, $3, $4, $5) + tail +) VALUES ($1, $2, $3, $4) RETURNING *; """ Apply this suggestion Suggestion importance[1-10]: 10 Why: The trailing comma in the SQL column list is a critical syntax error that would cause the query to fail. This is a high-priority fix as it affects core functionality.	10
	Fix mismatch between number of SQL parameters and columns in VALUES clause The VALUES clause in entry_relation_query has 5 placeholders ($1-$5) but only 4 columns are defined. Adjust the number of placeholders to match the columns. agents-api/agents_api/queries/entries/create_entries.py [46-54] entry_relation_query = """ INSERT INTO entry_relations ( session_id, head, relation, - tail, -) VALUES ($1, $2, $3, $4, $5) + tail +) VALUES ($1, $2, $3, $4) RETURNING *; """ Apply this suggestion Suggestion importance[1-10]: 10 Why: The mismatch between the number of SQL parameters ($1-$5) and columns (4) would cause a query execution error. This is a critical bug that needs immediate attention.	10
	Fix incorrect error message and status code for duplicate developer creation attempts The error message for UniqueViolationError incorrectly states "developer does not exist" when it should indicate a duplicate developer error. Update the error message to accurately reflect the conflict. agents-api/agents_api/queries/developers/create_developer.py [37-45] @rewrap_exceptions( { asyncpg.UniqueViolationError: partialclass( HTTPException, - status_code=404, - detail="The specified developer does not exist.", + status_code=409, + detail="A developer with this ID already exists.", ) } ) Apply this suggestion Suggestion importance[1-10]: 9 Why: The suggestion fixes a critical issue where a duplicate developer error is incorrectly reported as "developer not found" with a 404 status code. Using 409 Conflict with an accurate message is the correct HTTP semantic.	9
	Use function parameter instead of potentially None value from dictionary lookup In add_entry_relations function, the parameters list uses item.get() for session_id which could be None, but should use the session_id parameter passed to the function. agents-api/agents_api/queries/entries/create_entries.py [173-180] params.append( [ - item.get("session_id"), # $1 + session_id, # $1 item.get("head"), # $2 item.get("relation"), # $3 item.get("tail"), # $4 ] ) Apply this suggestion Suggestion importance[1-10]: 8 Why: Using item.get("session_id") could return None and cause database errors, while the session_id parameter is guaranteed to be valid. This is an important bug fix for data integrity.	8
	Prevent potential data loss when updating JSON fields by using COALESCE for safe merging The SQL query uses COALESCE with metadata and recall_options but directly overwrites them with || operator, which could lead to data loss - use COALESCE for these fields as well. agents-api/agents_api/queries/sessions/patch_session.py [20-25] -metadata = sessions.metadata || $5, -recall_options = sessions.recall_options || $10 +metadata = COALESCE(sessions.metadata || $5, sessions.metadata), +recall_options = COALESCE(sessions.recall_options || $10, sessions.recall_options) Apply this suggestion Suggestion importance[1-10]: 8 Why: The suggestion prevents potential data loss by safely merging JSON fields using COALESCE, ensuring that the original data is preserved if the update operation fails.	8
	Add validation for query timeout parameter to prevent invalid timeout values Add type validation for the query timeout parameter to ensure it's a positive number or None. agents-api/agents_api/queries/utils.py [91-93] +if query_timeout is not None and not isinstance(query_timeout, (int, float)) or isinstance(query_timeout, (int, float)) and query_timeout <= 0: + raise ValueError(f"Query timeout must be a positive number or None, got {query_timeout}") AsyncPGFetchArgs( query=query, args=variables, timeout=query_timeout ) Apply this suggestion Suggestion importance[1-10]: 7 Why: Important validation that prevents potential runtime errors and system instability by ensuring timeout values are valid before database operations.	7
	Add error handling for database pool initialization to prevent application crashes The app.state.postgres_pool access should be wrapped in try-except to handle AttributeError in case the state object hasn't been properly initialized. agents-api/agents_api/app.py [12-13] -if not app.state.postgres_pool: +try: + if not app.state.postgres_pool: + app.state.postgres_pool = await create_db_pool() +except AttributeError: app.state.postgres_pool = await create_db_pool() Apply this suggestion Suggestion importance[1-10]: 7 Why: The suggestion adds important error handling to prevent application crashes if the postgres_pool attribute is not properly initialized, improving application robustness.	7
Security	Add input validation to prevent SQL injection and database errors from empty or invalid query parameters Add validation for empty query strings and None values in variables list to prevent potential SQL injection or database errors. agents-api/agents_api/queries/utils.py [79-81] def prepare_pg_query_args( query_args: PGQueryArgs | list[PGQueryArgs], ) -> BatchedPreparedPGQueryArgs: + if not query_args: + raise ValueError("Query arguments cannot be empty") + if isinstance(query_args, tuple) and (not query_args[0] or None in query_args[1]): + raise ValueError("Query string cannot be empty and variables cannot contain None values") Apply this suggestion Suggestion importance[1-10]: 8 Why: This is a critical security enhancement that helps prevent SQL injection vulnerabilities and database errors by validating query inputs before execution.	8
General	Enhance error handling by providing more descriptive error messages for invalid query arguments Add error handling for invalid query arguments in prepare_pg_query_args(). The current implementation raises a generic ValueError which doesn't provide enough context about what was wrong with the arguments. agents-api/agents_api/queries/utils.py [114-115] case _: - raise ValueError("Invalid query arguments") + raise ValueError(f"Invalid query arguments. Expected tuple of (query, variables) or (query, variables, fetch_method). Got: {query_arg}") Apply this suggestion Suggestion importance[1-10]: 6 Why: The suggestion improves error diagnostics by providing more detailed error messages that include the actual invalid input, making it easier to debug issues with query arguments.	6