feat: Updated agents-api/agents_api/models/docs/se

agents-api/agents_api/models/docs/search_docs.py

@@ -1,3 +1,5 @@
+"""This module contains functions for searching documents in the CozoDB based on embedding queries."""
+
 from typing import Literal
 from uuid import UUID
 
@@ -7,6 +9,20 @@
 from ...clients.cozo import client
 
 
+"""
+Searches for document snippets in CozoDB by embedding query.
+
+Parameters:
+- owner_type (Literal["user", "agent"]): The type of the owner of the documents.
+- owner_id (UUID): The unique identifier of the owner.
+- query_embedding (list[float]): The embedding vector of the query.
+- k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3.
+- confidence (float, optional): The confidence threshold for filtering results. Defaults to 0.8.
+- client (CozoClient, optional): The CozoDB client instance. Defaults to a pre-configured client.
+
+Returns:
+- pd.DataFrame: A DataFrame containing the search results.
+"""
 def search_docs_snippets_by_embedding_query(
     owner_type: Literal["user", "agent"],
     owner_id: UUID,
@@ -16,8 +32,10 @@ def search_docs_snippets_by_embedding_query(
     client: CozoClient = client,
 ) -> pd.DataFrame:
     owner_id = str(owner_id)
+    # Calculate the search radius based on confidence level
     radius: float = 1.0 - confidence
 
+    # Construct the datalog query for searching document snippets
     query = f"""
     {{
         input[
@@ -59,6 +77,7 @@ def search_docs_snippets_by_embedding_query(
                 bind_vector: vector,
             }}
 
+        # Sort the results by distance to find the closest matches
         :sort distance
     }}"""