feat(agents-api): Add docs tests

Signed-off-by: Diwank Tomer <[email protected]>
julep-ai · Aug 10, 2024 · e0c3140 · e0c3140
1 parent 128f074
commit e0c3140
Show file tree

Hide file tree

Showing 19 changed files with 615 additions and 242 deletions.
diff --git a/agents-api/agents_api/common/utils/debug.py b/agents-api/agents_api/common/utils/debug.py
@@ -6,7 +6,9 @@ def pdb_on_exception(fn):
     def wrapper(*args, **kwargs):
         try:
             return fn(*args, **kwargs)
-        except Exception:
+        except Exception as exc:
+            print(repr(getattr(exc, "__cause__", exc)))
+
             import pdb
             import traceback
 

diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py
@@ -60,12 +60,13 @@ def create_agent(
 
     # Extract the agent data from the payload
     data.metadata = data.metadata or {}
+    data.default_settings = data.default_settings or {}
+
     data.instructions = (
         data.instructions
         if isinstance(data.instructions, list)
         else [data.instructions]
     )
-    data.default_settings = data.default_settings
 
     agent_data = data.model_dump(exclude_unset=True)
     default_settings = agent_data.pop("default_settings")

diff --git a/agents-api/agents_api/models/docs/create_doc.py b/agents-api/agents_api/models/docs/create_doc.py
@@ -30,7 +30,6 @@
     one=True,
     transform=lambda d: {
         "id": UUID(d["doc_id"]),
-        "content": [],  # <-- Note: we do not return content on creation
         **d,
     },
 )
@@ -87,15 +86,33 @@ def create_doc(
 
     create_snippets_query = f"""
         ?[{snippet_cols}] <- $snippet_rows
+
+        :create _snippets {{ {snippet_cols} }}
+        }} {{ 
+        ?[{snippet_cols}] <- $snippet_rows
         :insert snippets {{ {snippet_cols} }}
         :returning
     """
 
     # Construct the datalog query for creating the document and its snippets.
     create_doc_query = f"""
         ?[{doc_cols}] <- $doc_rows
+
+        :create _docs {{ {doc_cols} }}
+        }} {{
+        ?[{doc_cols}] <- $doc_rows
         :insert docs {{ {doc_cols} }}
         :returning
+        }} {{
+        snippet_rows[collect(content)] :=
+            *_snippets {{
+                content
+            }}
+
+        ?[{doc_cols}, content, created_at] :=
+            *_docs {{ {doc_cols} }},
+            snippet_rows[content],
+            created_at = now()
     """
 
     queries = [

diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py
@@ -19,6 +19,9 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, AssertionError)
+        and "Expected one result"
+        in repr(e): partialclass(HTTPException, status_code=404),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py
@@ -76,6 +76,24 @@ def search_docs_by_text(
                 doc_id
             }}
 
+        search_result[
+            doc_id,
+            snippet_data,
+            distance,
+        ] :=
+            input[owner_id, query],
+            candidate[doc_id],
+            ~snippets:lsh {{
+                doc_id,
+                index,
+                content
+                |
+                query: query,
+                k: {k},
+            }},
+            distance = 10000000,  # Very large distance to depict no distance
+            snippet_data = [index, content]
+
         search_result[
             doc_id,
             snippet_data,

diff --git a/agents-api/agents_api/routers/docs/__init__.py b/agents-api/agents_api/routers/docs/__init__.py
@@ -1,6 +1,6 @@
 # ruff: noqa: F401
 from .create_doc import create_agent_doc, create_user_doc
-from .delete_doc import delete_doc
+from .delete_doc import delete_agent_doc, delete_user_doc
 from .get_doc import get_doc
 from .list_docs import list_agent_docs, list_user_docs
 from .router import router

diff --git a/agents-api/agents_api/routers/docs/create_doc.py b/agents-api/agents_api/routers/docs/create_doc.py
@@ -23,7 +23,7 @@ async def create_user_doc(
         data=data,
     )
 
-    return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at)
+    return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at, jobs=[])
 
 
 @router.post("/agents/{agent_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"])
@@ -39,4 +39,4 @@ async def create_agent_doc(
         data=data,
     )
 
-    return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at)
+    return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at, jobs=[])
diff --git a/agents-api/agents_api/routers/docs/delete_doc.py b/agents-api/agents_api/routers/docs/delete_doc.py
@@ -10,8 +10,33 @@
 from .router import router
 
 
-@router.delete("/docs/{doc_id}", status_code=HTTP_202_ACCEPTED, tags=["docs"])
-async def delete_doc(
-    doc_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
+@router.delete(
+    "/agents/{agent_id}/docs/{doc_id}", status_code=HTTP_202_ACCEPTED, tags=["docs"]
+)
+async def delete_agent_doc(
+    doc_id: UUID4,
+    agent_id: UUID4,
+    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceDeletedResponse:
-    return delete_doc_query(developer_id=x_developer_id, doc_id=doc_id)
+    return delete_doc_query(
+        developer_id=x_developer_id,
+        owner_id=agent_id,
+        owner_type="agent",
+        doc_id=doc_id,
+    )
+
+
+@router.delete(
+    "/users/{user_id}/docs/{doc_id}", status_code=HTTP_202_ACCEPTED, tags=["docs"]
+)
+async def delete_user_doc(
+    doc_id: UUID4,
+    user_id: UUID4,
+    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+) -> ResourceDeletedResponse:
+    return delete_doc_query(
+        developer_id=x_developer_id,
+        owner_id=user_id,
+        owner_type="user",
+        doc_id=doc_id,
+    )
diff --git a/agents-api/agents_api/routers/docs/search_docs.py b/agents-api/agents_api/routers/docs/search_docs.py
@@ -74,6 +74,7 @@ async def search_user_docs(
         owner_id=user_id,
         **params,
     )
+
     end = time.time()
 
     time_taken = end - start
@@ -101,6 +102,7 @@ async def search_agent_docs(
         owner_id=agent_id,
         **params,
     )
+
     end = time.time()
 
     time_taken = end - start

diff --git a/agents-api/agents_api/routers/users/update_user.py b/agents-api/agents_api/routers/users/update_user.py
@@ -18,5 +18,5 @@ async def update_user(
     return update_user_query(
         developer_id=x_developer_id,
         user_id=user_id,
-        update_user=data,
+        data=data,
     )
diff --git a/agents-api/agents_api/web.py b/agents-api/agents_api/web.py
@@ -21,6 +21,7 @@
 from agents_api.exceptions import PromptTooBigError
 from agents_api.routers import (
     agents,
+    docs,
     jobs,
     sessions,
     tasks,
@@ -93,6 +94,7 @@ def register_exceptions(app: FastAPI):
 app.include_router(sessions.router)
 app.include_router(users.router)
 app.include_router(jobs.router)
+app.include_router(docs.router)
 app.include_router(tasks.router)
 
 

diff --git a/agents-api/migrations/migrate_1723307805_add_lsh_index_to_docs.py b/agents-api/migrations/migrate_1723307805_add_lsh_index_to_docs.py
@@ -0,0 +1,44 @@
+#/usr/bin/env python3
+
+MIGRATION_ID = "add_lsh_index_to_docs"
+CREATED_AT = 1723307805.007054
+
+# See: https://docs.cozodb.org/en/latest/vector.html#full-text-search-fts
+snippets_lsh_index = dict(
+    up="""
+    ::lsh create snippets:lsh {
+        extractor: content,
+        tokenizer: Simple,
+        filters: [Stopwords('en')],
+        n_perm: 200,
+        target_threshold: 0.9,
+        n_gram: 3,
+        false_positive_weight: 1.0,
+        false_negative_weight: 1.0,
+    }
+    """,
+    down="""
+    ::lsh drop snippets:lsh
+    """,
+)
+
+queries = [
+    snippets_lsh_index,
+]
+
+
+def run(client, queries):
+    joiner = "}\n\n{"
+
+    query = joiner.join(queries)
+    query = f"{{\n{query}\n}}"
+
+    client.run(query)
+
+
+def up(client):
+    run(client, [q["up"] for q in queries])
+
+
+def down(client):
+    run(client, [q["down"] for q in reversed(queries)])
diff --git a/agents-api/poetry.lock b/agents-api/poetry.lock
diff --git a/agents-api/tests/fixtures.py b/agents-api/tests/fixtures.py
@@ -157,6 +157,31 @@ def test_doc(
     )
 
 
+@fixture(scope="global")
+def test_user_doc(
+    client=cozo_client,
+    developer_id=test_developer_id,
+    user=test_user,
+):
+    doc = create_doc(
+        developer_id=developer_id,
+        owner_type="user",
+        owner_id=user.id,
+        data=CreateDocRequest(title="Hello", content=["World"]),
+        client=client,
+    )
+
+    yield doc
+
+    delete_doc(
+        developer_id=developer_id,
+        doc_id=doc.id,
+        owner_type="user",
+        owner_id=user.id,
+        client=client,
+    )
+
+
 @fixture(scope="global")
 def test_task(
     client=cozo_client,