-
Notifications
You must be signed in to change notification settings - Fork 811
/
list_docs.py
82 lines (67 loc) · 2.06 KB
/
list_docs.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
"""This module contains functions for querying document-related data from the 'cozodb' database using datalog queries."""
import json
from typing import Literal, Any
from uuid import UUID
from ..utils import cozo_query
@cozo_query
def ensure_owner_exists_query(
owner_type: Literal["user", "agent"],
owner_id: UUID,
) -> tuple[str, dict]:
owner_id = str(owner_id)
# Query to check if an owner (user or agent) exists in the database
query = f"""{{
# Convert owner_id to UUID and set as input
input[{owner_type}_id] <- [[to_uuid($owner_id)]]
# Retrieve owner_id if it exists in the database
?[
{owner_type}_id,
] := input[{owner_type}_id],
*{owner_type}s {{
{owner_type}_id,
}}
}}"""
return (query, {"owner_id": owner_id})
@cozo_query
def list_docs_snippets_by_owner_query(
owner_type: Literal["user", "agent"],
owner_id: UUID,
metadata_filter: dict[str, Any] = {},
) -> tuple[str, dict]:
owner_id = str(owner_id)
metadata_filter_str = ", ".join(
[
f"metadata->{json.dumps(k)} == {json.dumps(v)}"
for k, v in metadata_filter.items()
]
)
# Query to retrieve document snippets by owner (user or agent)
query = f"""
{{
# Convert owner_id to UUID and set as input
input[{owner_type}_id] <- [[to_uuid($owner_id)]]
# Retrieve documents and snippets associated with the owner
?[
{owner_type}_id,
doc_id,
title,
snippet,
snippet_idx,
created_at,
metadata,
] := input[{owner_type}_id],
*{owner_type}_docs {{
{owner_type}_id,
doc_id,
created_at,
metadata,
}},
{metadata_filter_str}
*information_snippets {{
doc_id,
snippet_idx,
title,
snippet,
}}
}}"""
return (query, {"owner_id": owner_id})