agno 2.0.1__py3-none-any.whl → 2.3.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- agno/agent/agent.py +6015 -2823
- agno/api/api.py +2 -0
- agno/api/os.py +1 -1
- agno/culture/__init__.py +3 -0
- agno/culture/manager.py +956 -0
- agno/db/async_postgres/__init__.py +3 -0
- agno/db/base.py +385 -6
- agno/db/dynamo/dynamo.py +388 -81
- agno/db/dynamo/schemas.py +47 -10
- agno/db/dynamo/utils.py +63 -4
- agno/db/firestore/firestore.py +435 -64
- agno/db/firestore/schemas.py +11 -0
- agno/db/firestore/utils.py +102 -4
- agno/db/gcs_json/gcs_json_db.py +384 -42
- agno/db/gcs_json/utils.py +60 -26
- agno/db/in_memory/in_memory_db.py +351 -66
- agno/db/in_memory/utils.py +60 -2
- agno/db/json/json_db.py +339 -48
- agno/db/json/utils.py +60 -26
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/v1_to_v2.py +510 -37
- agno/db/migrations/versions/__init__.py +0 -0
- agno/db/migrations/versions/v2_3_0.py +938 -0
- agno/db/mongo/__init__.py +15 -1
- agno/db/mongo/async_mongo.py +2036 -0
- agno/db/mongo/mongo.py +653 -76
- agno/db/mongo/schemas.py +13 -0
- agno/db/mongo/utils.py +80 -8
- agno/db/mysql/mysql.py +687 -25
- agno/db/mysql/schemas.py +61 -37
- agno/db/mysql/utils.py +60 -2
- agno/db/postgres/__init__.py +2 -1
- agno/db/postgres/async_postgres.py +2001 -0
- agno/db/postgres/postgres.py +676 -57
- agno/db/postgres/schemas.py +43 -18
- agno/db/postgres/utils.py +164 -2
- agno/db/redis/redis.py +344 -38
- agno/db/redis/schemas.py +18 -0
- agno/db/redis/utils.py +60 -2
- agno/db/schemas/__init__.py +2 -1
- agno/db/schemas/culture.py +120 -0
- agno/db/schemas/memory.py +13 -0
- agno/db/singlestore/schemas.py +26 -1
- agno/db/singlestore/singlestore.py +687 -53
- agno/db/singlestore/utils.py +60 -2
- agno/db/sqlite/__init__.py +2 -1
- agno/db/sqlite/async_sqlite.py +2371 -0
- agno/db/sqlite/schemas.py +24 -0
- agno/db/sqlite/sqlite.py +774 -85
- agno/db/sqlite/utils.py +168 -5
- agno/db/surrealdb/__init__.py +3 -0
- agno/db/surrealdb/metrics.py +292 -0
- agno/db/surrealdb/models.py +309 -0
- agno/db/surrealdb/queries.py +71 -0
- agno/db/surrealdb/surrealdb.py +1361 -0
- agno/db/surrealdb/utils.py +147 -0
- agno/db/utils.py +50 -22
- agno/eval/accuracy.py +50 -43
- agno/eval/performance.py +6 -3
- agno/eval/reliability.py +6 -3
- agno/eval/utils.py +33 -16
- agno/exceptions.py +68 -1
- agno/filters.py +354 -0
- agno/guardrails/__init__.py +6 -0
- agno/guardrails/base.py +19 -0
- agno/guardrails/openai.py +144 -0
- agno/guardrails/pii.py +94 -0
- agno/guardrails/prompt_injection.py +52 -0
- agno/integrations/discord/client.py +1 -0
- agno/knowledge/chunking/agentic.py +13 -10
- agno/knowledge/chunking/fixed.py +1 -1
- agno/knowledge/chunking/semantic.py +40 -8
- agno/knowledge/chunking/strategy.py +59 -15
- agno/knowledge/embedder/aws_bedrock.py +9 -4
- agno/knowledge/embedder/azure_openai.py +54 -0
- agno/knowledge/embedder/base.py +2 -0
- agno/knowledge/embedder/cohere.py +184 -5
- agno/knowledge/embedder/fastembed.py +1 -1
- agno/knowledge/embedder/google.py +79 -1
- agno/knowledge/embedder/huggingface.py +9 -4
- agno/knowledge/embedder/jina.py +63 -0
- agno/knowledge/embedder/mistral.py +78 -11
- agno/knowledge/embedder/nebius.py +1 -1
- agno/knowledge/embedder/ollama.py +13 -0
- agno/knowledge/embedder/openai.py +37 -65
- agno/knowledge/embedder/sentence_transformer.py +8 -4
- agno/knowledge/embedder/vllm.py +262 -0
- agno/knowledge/embedder/voyageai.py +69 -16
- agno/knowledge/knowledge.py +594 -186
- agno/knowledge/reader/base.py +9 -2
- agno/knowledge/reader/csv_reader.py +8 -10
- agno/knowledge/reader/docx_reader.py +5 -6
- agno/knowledge/reader/field_labeled_csv_reader.py +290 -0
- agno/knowledge/reader/json_reader.py +6 -5
- agno/knowledge/reader/markdown_reader.py +13 -13
- agno/knowledge/reader/pdf_reader.py +43 -68
- agno/knowledge/reader/pptx_reader.py +101 -0
- agno/knowledge/reader/reader_factory.py +51 -6
- agno/knowledge/reader/s3_reader.py +3 -15
- agno/knowledge/reader/tavily_reader.py +194 -0
- agno/knowledge/reader/text_reader.py +13 -13
- agno/knowledge/reader/web_search_reader.py +2 -43
- agno/knowledge/reader/website_reader.py +43 -25
- agno/knowledge/reranker/__init__.py +2 -8
- agno/knowledge/types.py +9 -0
- agno/knowledge/utils.py +20 -0
- agno/media.py +72 -0
- agno/memory/manager.py +336 -82
- agno/models/aimlapi/aimlapi.py +2 -2
- agno/models/anthropic/claude.py +183 -37
- agno/models/aws/bedrock.py +52 -112
- agno/models/aws/claude.py +33 -1
- agno/models/azure/ai_foundry.py +33 -15
- agno/models/azure/openai_chat.py +25 -8
- agno/models/base.py +999 -519
- agno/models/cerebras/cerebras.py +19 -13
- agno/models/cerebras/cerebras_openai.py +8 -5
- agno/models/cohere/chat.py +27 -1
- agno/models/cometapi/__init__.py +5 -0
- agno/models/cometapi/cometapi.py +57 -0
- agno/models/dashscope/dashscope.py +1 -0
- agno/models/deepinfra/deepinfra.py +2 -2
- agno/models/deepseek/deepseek.py +2 -2
- agno/models/fireworks/fireworks.py +2 -2
- agno/models/google/gemini.py +103 -31
- agno/models/groq/groq.py +28 -11
- agno/models/huggingface/huggingface.py +2 -1
- agno/models/internlm/internlm.py +2 -2
- agno/models/langdb/langdb.py +4 -4
- agno/models/litellm/chat.py +18 -1
- agno/models/litellm/litellm_openai.py +2 -2
- agno/models/llama_cpp/__init__.py +5 -0
- agno/models/llama_cpp/llama_cpp.py +22 -0
- agno/models/message.py +139 -0
- agno/models/meta/llama.py +27 -10
- agno/models/meta/llama_openai.py +5 -17
- agno/models/nebius/nebius.py +6 -6
- agno/models/nexus/__init__.py +3 -0
- agno/models/nexus/nexus.py +22 -0
- agno/models/nvidia/nvidia.py +2 -2
- agno/models/ollama/chat.py +59 -5
- agno/models/openai/chat.py +69 -29
- agno/models/openai/responses.py +103 -106
- agno/models/openrouter/openrouter.py +41 -3
- agno/models/perplexity/perplexity.py +4 -5
- agno/models/portkey/portkey.py +3 -3
- agno/models/requesty/__init__.py +5 -0
- agno/models/requesty/requesty.py +52 -0
- agno/models/response.py +77 -1
- agno/models/sambanova/sambanova.py +2 -2
- agno/models/siliconflow/__init__.py +5 -0
- agno/models/siliconflow/siliconflow.py +25 -0
- agno/models/together/together.py +2 -2
- agno/models/utils.py +254 -8
- agno/models/vercel/v0.py +2 -2
- agno/models/vertexai/__init__.py +0 -0
- agno/models/vertexai/claude.py +96 -0
- agno/models/vllm/vllm.py +1 -0
- agno/models/xai/xai.py +3 -2
- agno/os/app.py +543 -178
- agno/os/auth.py +24 -14
- agno/os/config.py +1 -0
- agno/os/interfaces/__init__.py +1 -0
- agno/os/interfaces/a2a/__init__.py +3 -0
- agno/os/interfaces/a2a/a2a.py +42 -0
- agno/os/interfaces/a2a/router.py +250 -0
- agno/os/interfaces/a2a/utils.py +924 -0
- agno/os/interfaces/agui/agui.py +23 -7
- agno/os/interfaces/agui/router.py +27 -3
- agno/os/interfaces/agui/utils.py +242 -142
- agno/os/interfaces/base.py +6 -2
- agno/os/interfaces/slack/router.py +81 -23
- agno/os/interfaces/slack/slack.py +29 -14
- agno/os/interfaces/whatsapp/router.py +11 -4
- agno/os/interfaces/whatsapp/whatsapp.py +14 -7
- agno/os/mcp.py +111 -54
- agno/os/middleware/__init__.py +7 -0
- agno/os/middleware/jwt.py +233 -0
- agno/os/router.py +556 -139
- agno/os/routers/evals/evals.py +71 -34
- agno/os/routers/evals/schemas.py +31 -31
- agno/os/routers/evals/utils.py +6 -5
- agno/os/routers/health.py +31 -0
- agno/os/routers/home.py +52 -0
- agno/os/routers/knowledge/knowledge.py +185 -38
- agno/os/routers/knowledge/schemas.py +82 -22
- agno/os/routers/memory/memory.py +158 -53
- agno/os/routers/memory/schemas.py +20 -16
- agno/os/routers/metrics/metrics.py +20 -8
- agno/os/routers/metrics/schemas.py +16 -16
- agno/os/routers/session/session.py +499 -38
- agno/os/schema.py +308 -198
- agno/os/utils.py +401 -41
- agno/reasoning/anthropic.py +80 -0
- agno/reasoning/azure_ai_foundry.py +2 -2
- agno/reasoning/deepseek.py +2 -2
- agno/reasoning/default.py +3 -1
- agno/reasoning/gemini.py +73 -0
- agno/reasoning/groq.py +2 -2
- agno/reasoning/ollama.py +2 -2
- agno/reasoning/openai.py +7 -2
- agno/reasoning/vertexai.py +76 -0
- agno/run/__init__.py +6 -0
- agno/run/agent.py +248 -94
- agno/run/base.py +44 -5
- agno/run/team.py +238 -97
- agno/run/workflow.py +144 -33
- agno/session/agent.py +105 -89
- agno/session/summary.py +65 -25
- agno/session/team.py +176 -96
- agno/session/workflow.py +406 -40
- agno/team/team.py +3854 -1610
- agno/tools/dalle.py +2 -4
- agno/tools/decorator.py +4 -2
- agno/tools/duckduckgo.py +15 -11
- agno/tools/e2b.py +14 -7
- agno/tools/eleven_labs.py +23 -25
- agno/tools/exa.py +21 -16
- agno/tools/file.py +153 -23
- agno/tools/file_generation.py +350 -0
- agno/tools/firecrawl.py +4 -4
- agno/tools/function.py +250 -30
- agno/tools/gmail.py +238 -14
- agno/tools/google_drive.py +270 -0
- agno/tools/googlecalendar.py +36 -8
- agno/tools/googlesheets.py +20 -5
- agno/tools/jira.py +20 -0
- agno/tools/knowledge.py +3 -3
- agno/tools/mcp/__init__.py +10 -0
- agno/tools/mcp/mcp.py +331 -0
- agno/tools/mcp/multi_mcp.py +347 -0
- agno/tools/mcp/params.py +24 -0
- agno/tools/mcp_toolbox.py +284 -0
- agno/tools/mem0.py +11 -17
- agno/tools/memori.py +1 -53
- agno/tools/memory.py +419 -0
- agno/tools/models/nebius.py +5 -5
- agno/tools/models_labs.py +20 -10
- agno/tools/notion.py +204 -0
- agno/tools/parallel.py +314 -0
- agno/tools/scrapegraph.py +58 -31
- agno/tools/searxng.py +2 -2
- agno/tools/serper.py +2 -2
- agno/tools/slack.py +18 -3
- agno/tools/spider.py +2 -2
- agno/tools/tavily.py +146 -0
- agno/tools/whatsapp.py +1 -1
- agno/tools/workflow.py +278 -0
- agno/tools/yfinance.py +12 -11
- agno/utils/agent.py +820 -0
- agno/utils/audio.py +27 -0
- agno/utils/common.py +90 -1
- agno/utils/events.py +217 -2
- agno/utils/gemini.py +180 -22
- agno/utils/hooks.py +57 -0
- agno/utils/http.py +111 -0
- agno/utils/knowledge.py +12 -5
- agno/utils/log.py +1 -0
- agno/utils/mcp.py +92 -2
- agno/utils/media.py +188 -10
- agno/utils/merge_dict.py +22 -1
- agno/utils/message.py +60 -0
- agno/utils/models/claude.py +40 -11
- agno/utils/print_response/agent.py +105 -21
- agno/utils/print_response/team.py +103 -38
- agno/utils/print_response/workflow.py +251 -34
- agno/utils/reasoning.py +22 -1
- agno/utils/serialize.py +32 -0
- agno/utils/streamlit.py +16 -10
- agno/utils/string.py +41 -0
- agno/utils/team.py +98 -9
- agno/utils/tools.py +1 -1
- agno/vectordb/base.py +23 -4
- agno/vectordb/cassandra/cassandra.py +65 -9
- agno/vectordb/chroma/chromadb.py +182 -38
- agno/vectordb/clickhouse/clickhousedb.py +64 -11
- agno/vectordb/couchbase/couchbase.py +105 -10
- agno/vectordb/lancedb/lance_db.py +124 -133
- agno/vectordb/langchaindb/langchaindb.py +25 -7
- agno/vectordb/lightrag/lightrag.py +17 -3
- agno/vectordb/llamaindex/__init__.py +3 -0
- agno/vectordb/llamaindex/llamaindexdb.py +46 -7
- agno/vectordb/milvus/milvus.py +126 -9
- agno/vectordb/mongodb/__init__.py +7 -1
- agno/vectordb/mongodb/mongodb.py +112 -7
- agno/vectordb/pgvector/pgvector.py +142 -21
- agno/vectordb/pineconedb/pineconedb.py +80 -8
- agno/vectordb/qdrant/qdrant.py +125 -39
- agno/vectordb/redis/__init__.py +9 -0
- agno/vectordb/redis/redisdb.py +694 -0
- agno/vectordb/singlestore/singlestore.py +111 -25
- agno/vectordb/surrealdb/surrealdb.py +31 -5
- agno/vectordb/upstashdb/upstashdb.py +76 -8
- agno/vectordb/weaviate/weaviate.py +86 -15
- agno/workflow/__init__.py +2 -0
- agno/workflow/agent.py +299 -0
- agno/workflow/condition.py +112 -18
- agno/workflow/loop.py +69 -10
- agno/workflow/parallel.py +266 -118
- agno/workflow/router.py +110 -17
- agno/workflow/step.py +638 -129
- agno/workflow/steps.py +65 -6
- agno/workflow/types.py +61 -23
- agno/workflow/workflow.py +2085 -272
- {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/METADATA +182 -58
- agno-2.3.0.dist-info/RECORD +577 -0
- agno/knowledge/reader/url_reader.py +0 -128
- agno/tools/googlesearch.py +0 -98
- agno/tools/mcp.py +0 -610
- agno/utils/models/aws_claude.py +0 -170
- agno-2.0.1.dist-info/RECORD +0 -515
- {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/WHEEL +0 -0
- {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/licenses/LICENSE +0 -0
- {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/top_level.txt +0 -0
agno/db/firestore/firestore.py
CHANGED
|
@@ -6,19 +6,25 @@ from uuid import uuid4
|
|
|
6
6
|
from agno.db.base import BaseDb, SessionType
|
|
7
7
|
from agno.db.firestore.utils import (
|
|
8
8
|
apply_pagination,
|
|
9
|
+
apply_pagination_to_records,
|
|
9
10
|
apply_sorting,
|
|
11
|
+
apply_sorting_to_records,
|
|
10
12
|
bulk_upsert_metrics,
|
|
11
13
|
calculate_date_metrics,
|
|
12
14
|
create_collection_indexes,
|
|
15
|
+
deserialize_cultural_knowledge_from_db,
|
|
13
16
|
fetch_all_sessions_data,
|
|
14
17
|
get_dates_to_calculate_metrics_for,
|
|
18
|
+
serialize_cultural_knowledge_for_db,
|
|
15
19
|
)
|
|
20
|
+
from agno.db.schemas.culture import CulturalKnowledge
|
|
16
21
|
from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
|
|
17
22
|
from agno.db.schemas.knowledge import KnowledgeRow
|
|
18
23
|
from agno.db.schemas.memory import UserMemory
|
|
19
24
|
from agno.db.utils import deserialize_session_json_fields, serialize_session_json_fields
|
|
20
25
|
from agno.session import AgentSession, Session, TeamSession, WorkflowSession
|
|
21
26
|
from agno.utils.log import log_debug, log_error, log_info
|
|
27
|
+
from agno.utils.string import generate_id
|
|
22
28
|
|
|
23
29
|
try:
|
|
24
30
|
from google.cloud.firestore import Client, FieldFilter # type: ignore[import-untyped]
|
|
@@ -38,6 +44,8 @@ class FirestoreDb(BaseDb):
|
|
|
38
44
|
metrics_collection: Optional[str] = None,
|
|
39
45
|
eval_collection: Optional[str] = None,
|
|
40
46
|
knowledge_collection: Optional[str] = None,
|
|
47
|
+
culture_collection: Optional[str] = None,
|
|
48
|
+
id: Optional[str] = None,
|
|
41
49
|
):
|
|
42
50
|
"""
|
|
43
51
|
Interface for interacting with a Firestore database.
|
|
@@ -50,16 +58,24 @@ class FirestoreDb(BaseDb):
|
|
|
50
58
|
metrics_collection (Optional[str]): Name of the collection to store metrics.
|
|
51
59
|
eval_collection (Optional[str]): Name of the collection to store evaluation runs.
|
|
52
60
|
knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
|
|
61
|
+
culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
|
|
62
|
+
id (Optional[str]): ID of the database.
|
|
53
63
|
|
|
54
64
|
Raises:
|
|
55
65
|
ValueError: If neither project_id nor db_client is provided.
|
|
56
66
|
"""
|
|
67
|
+
if id is None:
|
|
68
|
+
seed = project_id or str(db_client)
|
|
69
|
+
id = generate_id(seed)
|
|
70
|
+
|
|
57
71
|
super().__init__(
|
|
72
|
+
id=id,
|
|
58
73
|
session_table=session_collection,
|
|
59
74
|
memory_table=memory_collection,
|
|
60
75
|
metrics_table=metrics_collection,
|
|
61
76
|
eval_table=eval_collection,
|
|
62
77
|
knowledge_table=knowledge_collection,
|
|
78
|
+
culture_table=culture_collection,
|
|
63
79
|
)
|
|
64
80
|
|
|
65
81
|
_client: Optional[Client] = db_client
|
|
@@ -73,6 +89,17 @@ class FirestoreDb(BaseDb):
|
|
|
73
89
|
|
|
74
90
|
# -- DB methods --
|
|
75
91
|
|
|
92
|
+
def table_exists(self, table_name: str) -> bool:
|
|
93
|
+
"""Check if a collection with the given name exists in the Firestore database.
|
|
94
|
+
|
|
95
|
+
Args:
|
|
96
|
+
table_name: Name of the collection to check
|
|
97
|
+
|
|
98
|
+
Returns:
|
|
99
|
+
bool: True if the collection exists in the database, False otherwise
|
|
100
|
+
"""
|
|
101
|
+
return table_name in self.db_client.list_collections()
|
|
102
|
+
|
|
76
103
|
def _get_collection(self, table_type: str, create_collection_if_not_found: Optional[bool] = True):
|
|
77
104
|
"""Get or create a collection based on table type.
|
|
78
105
|
|
|
@@ -138,6 +165,17 @@ class FirestoreDb(BaseDb):
|
|
|
138
165
|
)
|
|
139
166
|
return self.knowledge_collection
|
|
140
167
|
|
|
168
|
+
if table_type == "culture":
|
|
169
|
+
if not hasattr(self, "culture_collection"):
|
|
170
|
+
if self.culture_table_name is None:
|
|
171
|
+
raise ValueError("Culture collection was not provided on initialization")
|
|
172
|
+
self.culture_collection = self._get_or_create_collection(
|
|
173
|
+
collection_name=self.culture_table_name,
|
|
174
|
+
collection_type="culture",
|
|
175
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
176
|
+
)
|
|
177
|
+
return self.culture_collection
|
|
178
|
+
|
|
141
179
|
raise ValueError(f"Unknown table type: {table_type}")
|
|
142
180
|
|
|
143
181
|
def _get_or_create_collection(
|
|
@@ -197,7 +235,15 @@ class FirestoreDb(BaseDb):
|
|
|
197
235
|
|
|
198
236
|
except Exception as e:
|
|
199
237
|
log_error(f"Error deleting session: {e}")
|
|
200
|
-
|
|
238
|
+
raise e
|
|
239
|
+
|
|
240
|
+
def get_latest_schema_version(self):
|
|
241
|
+
"""Get the latest version of the database schema."""
|
|
242
|
+
pass
|
|
243
|
+
|
|
244
|
+
def upsert_schema_version(self, version: str) -> None:
|
|
245
|
+
"""Upsert the schema version into the database."""
|
|
246
|
+
pass
|
|
201
247
|
|
|
202
248
|
def delete_sessions(self, session_ids: List[str]) -> None:
|
|
203
249
|
"""Delete multiple sessions from the database.
|
|
@@ -222,6 +268,7 @@ class FirestoreDb(BaseDb):
|
|
|
222
268
|
|
|
223
269
|
except Exception as e:
|
|
224
270
|
log_error(f"Error deleting sessions: {e}")
|
|
271
|
+
raise e
|
|
225
272
|
|
|
226
273
|
def get_session(
|
|
227
274
|
self,
|
|
@@ -234,8 +281,8 @@ class FirestoreDb(BaseDb):
|
|
|
234
281
|
|
|
235
282
|
Args:
|
|
236
283
|
session_id (str): The ID of the session to get.
|
|
284
|
+
session_type (SessionType): The type of session to get.
|
|
237
285
|
user_id (Optional[str]): The ID of the user to get the session for.
|
|
238
|
-
session_type (Optional[SessionType]): The type of session to get.
|
|
239
286
|
deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
|
|
240
287
|
|
|
241
288
|
Returns:
|
|
@@ -252,8 +299,6 @@ class FirestoreDb(BaseDb):
|
|
|
252
299
|
|
|
253
300
|
if user_id is not None:
|
|
254
301
|
query = query.where(filter=FieldFilter("user_id", "==", user_id))
|
|
255
|
-
if session_type is not None:
|
|
256
|
-
query = query.where(filter=FieldFilter("session_type", "==", session_type.value))
|
|
257
302
|
|
|
258
303
|
docs = query.stream()
|
|
259
304
|
result = None
|
|
@@ -273,12 +318,14 @@ class FirestoreDb(BaseDb):
|
|
|
273
318
|
return AgentSession.from_dict(session)
|
|
274
319
|
elif session_type == SessionType.TEAM:
|
|
275
320
|
return TeamSession.from_dict(session)
|
|
276
|
-
|
|
321
|
+
elif session_type == SessionType.WORKFLOW:
|
|
277
322
|
return WorkflowSession.from_dict(session)
|
|
323
|
+
else:
|
|
324
|
+
raise ValueError(f"Invalid session type: {session_type}")
|
|
278
325
|
|
|
279
326
|
except Exception as e:
|
|
280
327
|
log_error(f"Exception reading session: {e}")
|
|
281
|
-
|
|
328
|
+
raise e
|
|
282
329
|
|
|
283
330
|
def get_sessions(
|
|
284
331
|
self,
|
|
@@ -392,7 +439,7 @@ class FirestoreDb(BaseDb):
|
|
|
392
439
|
|
|
393
440
|
except Exception as e:
|
|
394
441
|
log_error(f"Exception reading sessions: {e}")
|
|
395
|
-
|
|
442
|
+
raise e
|
|
396
443
|
|
|
397
444
|
def rename_session(
|
|
398
445
|
self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
|
|
@@ -447,7 +494,7 @@ class FirestoreDb(BaseDb):
|
|
|
447
494
|
|
|
448
495
|
except Exception as e:
|
|
449
496
|
log_error(f"Exception renaming session: {e}")
|
|
450
|
-
|
|
497
|
+
raise e
|
|
451
498
|
|
|
452
499
|
def upsert_session(
|
|
453
500
|
self, session: Session, deserialize: Optional[bool] = True
|
|
@@ -544,15 +591,53 @@ class FirestoreDb(BaseDb):
|
|
|
544
591
|
|
|
545
592
|
except Exception as e:
|
|
546
593
|
log_error(f"Exception upserting session: {e}")
|
|
547
|
-
|
|
594
|
+
raise e
|
|
595
|
+
|
|
596
|
+
def upsert_sessions(
|
|
597
|
+
self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
|
|
598
|
+
) -> List[Union[Session, Dict[str, Any]]]:
|
|
599
|
+
"""
|
|
600
|
+
Bulk upsert multiple sessions for improved performance on large datasets.
|
|
601
|
+
|
|
602
|
+
Args:
|
|
603
|
+
sessions (List[Session]): List of sessions to upsert.
|
|
604
|
+
deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
|
|
605
|
+
|
|
606
|
+
Returns:
|
|
607
|
+
List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
|
|
608
|
+
|
|
609
|
+
Raises:
|
|
610
|
+
Exception: If an error occurs during bulk upsert.
|
|
611
|
+
"""
|
|
612
|
+
if not sessions:
|
|
613
|
+
return []
|
|
614
|
+
|
|
615
|
+
try:
|
|
616
|
+
log_info(
|
|
617
|
+
f"FirestoreDb doesn't support efficient bulk operations, falling back to individual upserts for {len(sessions)} sessions"
|
|
618
|
+
)
|
|
619
|
+
|
|
620
|
+
# Fall back to individual upserts
|
|
621
|
+
results = []
|
|
622
|
+
for session in sessions:
|
|
623
|
+
if session is not None:
|
|
624
|
+
result = self.upsert_session(session, deserialize=deserialize)
|
|
625
|
+
if result is not None:
|
|
626
|
+
results.append(result)
|
|
627
|
+
return results
|
|
628
|
+
|
|
629
|
+
except Exception as e:
|
|
630
|
+
log_error(f"Exception during bulk session upsert: {e}")
|
|
631
|
+
return []
|
|
548
632
|
|
|
549
633
|
# -- Memory methods --
|
|
550
634
|
|
|
551
|
-
def delete_user_memory(self, memory_id: str):
|
|
635
|
+
def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None):
|
|
552
636
|
"""Delete a user memory from the database.
|
|
553
637
|
|
|
554
638
|
Args:
|
|
555
639
|
memory_id (str): The ID of the memory to delete.
|
|
640
|
+
user_id (Optional[str]): The ID of the user (optional, for filtering).
|
|
556
641
|
|
|
557
642
|
Returns:
|
|
558
643
|
bool: True if the memory was deleted, False otherwise.
|
|
@@ -562,27 +647,41 @@ class FirestoreDb(BaseDb):
|
|
|
562
647
|
"""
|
|
563
648
|
try:
|
|
564
649
|
collection_ref = self._get_collection(table_type="memories")
|
|
565
|
-
docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
|
|
566
|
-
|
|
567
|
-
deleted_count = 0
|
|
568
|
-
for doc in docs:
|
|
569
|
-
doc.reference.delete()
|
|
570
|
-
deleted_count += 1
|
|
571
650
|
|
|
572
|
-
|
|
573
|
-
if
|
|
574
|
-
|
|
651
|
+
# If user_id is provided, verify the memory belongs to the user before deleting
|
|
652
|
+
if user_id:
|
|
653
|
+
docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
|
|
654
|
+
for doc in docs:
|
|
655
|
+
data = doc.to_dict()
|
|
656
|
+
if data.get("user_id") != user_id:
|
|
657
|
+
log_debug(f"Memory {memory_id} does not belong to user {user_id}")
|
|
658
|
+
return
|
|
659
|
+
doc.reference.delete()
|
|
660
|
+
log_debug(f"Successfully deleted user memory id: {memory_id}")
|
|
661
|
+
return
|
|
575
662
|
else:
|
|
576
|
-
|
|
663
|
+
docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
|
|
664
|
+
deleted_count = 0
|
|
665
|
+
for doc in docs:
|
|
666
|
+
doc.reference.delete()
|
|
667
|
+
deleted_count += 1
|
|
668
|
+
|
|
669
|
+
success = deleted_count > 0
|
|
670
|
+
if success:
|
|
671
|
+
log_debug(f"Successfully deleted user memory id: {memory_id}")
|
|
672
|
+
else:
|
|
673
|
+
log_debug(f"No user memory found with id: {memory_id}")
|
|
577
674
|
|
|
578
675
|
except Exception as e:
|
|
579
676
|
log_error(f"Error deleting user memory: {e}")
|
|
677
|
+
raise e
|
|
580
678
|
|
|
581
|
-
def delete_user_memories(self, memory_ids: List[str]) -> None:
|
|
679
|
+
def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
|
|
582
680
|
"""Delete user memories from the database.
|
|
583
681
|
|
|
584
682
|
Args:
|
|
585
683
|
memory_ids (List[str]): The IDs of the memories to delete.
|
|
684
|
+
user_id (Optional[str]): The ID of the user (optional, for filtering).
|
|
586
685
|
|
|
587
686
|
Raises:
|
|
588
687
|
Exception: If there is an error deleting the memories.
|
|
@@ -592,11 +691,21 @@ class FirestoreDb(BaseDb):
|
|
|
592
691
|
batch = self.db_client.batch()
|
|
593
692
|
deleted_count = 0
|
|
594
693
|
|
|
595
|
-
|
|
596
|
-
|
|
597
|
-
for
|
|
598
|
-
|
|
599
|
-
|
|
694
|
+
# If user_id is provided, filter memory_ids to only those belonging to the user
|
|
695
|
+
if user_id:
|
|
696
|
+
for memory_id in memory_ids:
|
|
697
|
+
docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
|
|
698
|
+
for doc in docs:
|
|
699
|
+
data = doc.to_dict()
|
|
700
|
+
if data.get("user_id") == user_id:
|
|
701
|
+
batch.delete(doc.reference)
|
|
702
|
+
deleted_count += 1
|
|
703
|
+
else:
|
|
704
|
+
for memory_id in memory_ids:
|
|
705
|
+
docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
|
|
706
|
+
for doc in docs:
|
|
707
|
+
batch.delete(doc.reference)
|
|
708
|
+
deleted_count += 1
|
|
600
709
|
|
|
601
710
|
batch.commit()
|
|
602
711
|
|
|
@@ -607,6 +716,7 @@ class FirestoreDb(BaseDb):
|
|
|
607
716
|
|
|
608
717
|
except Exception as e:
|
|
609
718
|
log_error(f"Error deleting memories: {e}")
|
|
719
|
+
raise e
|
|
610
720
|
|
|
611
721
|
def get_all_memory_topics(self, create_collection_if_not_found: Optional[bool] = True) -> List[str]:
|
|
612
722
|
"""Get all memory topics from the database.
|
|
@@ -634,15 +744,18 @@ class FirestoreDb(BaseDb):
|
|
|
634
744
|
return [topic for topic in all_topics if topic]
|
|
635
745
|
|
|
636
746
|
except Exception as e:
|
|
637
|
-
log_error(f"Exception
|
|
638
|
-
|
|
747
|
+
log_error(f"Exception getting all memory topics: {e}")
|
|
748
|
+
raise e
|
|
639
749
|
|
|
640
|
-
def get_user_memory(
|
|
750
|
+
def get_user_memory(
|
|
751
|
+
self, memory_id: str, deserialize: Optional[bool] = True, user_id: Optional[str] = None
|
|
752
|
+
) -> Optional[UserMemory]:
|
|
641
753
|
"""Get a memory from the database.
|
|
642
754
|
|
|
643
755
|
Args:
|
|
644
756
|
memory_id (str): The ID of the memory to get.
|
|
645
757
|
deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
|
|
758
|
+
user_id (Optional[str]): The ID of the user (optional, for filtering).
|
|
646
759
|
|
|
647
760
|
Returns:
|
|
648
761
|
Optional[UserMemory]:
|
|
@@ -661,14 +774,21 @@ class FirestoreDb(BaseDb):
|
|
|
661
774
|
result = doc.to_dict()
|
|
662
775
|
break
|
|
663
776
|
|
|
664
|
-
if result is None
|
|
777
|
+
if result is None:
|
|
778
|
+
return None
|
|
779
|
+
|
|
780
|
+
# Filter by user_id if provided
|
|
781
|
+
if user_id and result.get("user_id") != user_id:
|
|
782
|
+
return None
|
|
783
|
+
|
|
784
|
+
if not deserialize:
|
|
665
785
|
return result
|
|
666
786
|
|
|
667
787
|
return UserMemory.from_dict(result)
|
|
668
788
|
|
|
669
789
|
except Exception as e:
|
|
670
|
-
log_error(f"Exception
|
|
671
|
-
|
|
790
|
+
log_error(f"Exception getting user memory: {e}")
|
|
791
|
+
raise e
|
|
672
792
|
|
|
673
793
|
def get_user_memories(
|
|
674
794
|
self,
|
|
@@ -746,8 +866,8 @@ class FirestoreDb(BaseDb):
|
|
|
746
866
|
return [UserMemory.from_dict(record) for record in records]
|
|
747
867
|
|
|
748
868
|
except Exception as e:
|
|
749
|
-
log_error(f"Exception
|
|
750
|
-
|
|
869
|
+
log_error(f"Exception getting user memories: {e}")
|
|
870
|
+
raise e
|
|
751
871
|
|
|
752
872
|
def get_user_memory_stats(
|
|
753
873
|
self,
|
|
@@ -768,23 +888,26 @@ class FirestoreDb(BaseDb):
|
|
|
768
888
|
"""
|
|
769
889
|
try:
|
|
770
890
|
collection_ref = self._get_collection(table_type="memories")
|
|
771
|
-
|
|
891
|
+
|
|
892
|
+
query = collection_ref.where(filter=FieldFilter("user_id", "!=", None))
|
|
893
|
+
|
|
894
|
+
docs = query.stream()
|
|
772
895
|
|
|
773
896
|
user_stats = {}
|
|
774
897
|
for doc in docs:
|
|
775
898
|
data = doc.to_dict()
|
|
776
|
-
|
|
777
|
-
if
|
|
778
|
-
if
|
|
779
|
-
user_stats[
|
|
780
|
-
"user_id":
|
|
899
|
+
current_user_id = data.get("user_id")
|
|
900
|
+
if current_user_id:
|
|
901
|
+
if current_user_id not in user_stats:
|
|
902
|
+
user_stats[current_user_id] = {
|
|
903
|
+
"user_id": current_user_id,
|
|
781
904
|
"total_memories": 0,
|
|
782
905
|
"last_memory_updated_at": 0,
|
|
783
906
|
}
|
|
784
|
-
user_stats[
|
|
907
|
+
user_stats[current_user_id]["total_memories"] += 1
|
|
785
908
|
updated_at = data.get("updated_at", 0)
|
|
786
|
-
if updated_at > user_stats[
|
|
787
|
-
user_stats[
|
|
909
|
+
if updated_at > user_stats[current_user_id]["last_memory_updated_at"]:
|
|
910
|
+
user_stats[current_user_id]["last_memory_updated_at"] = updated_at
|
|
788
911
|
|
|
789
912
|
# Convert to list and sort
|
|
790
913
|
formatted_results = list(user_stats.values())
|
|
@@ -803,7 +926,7 @@ class FirestoreDb(BaseDb):
|
|
|
803
926
|
|
|
804
927
|
except Exception as e:
|
|
805
928
|
log_error(f"Exception getting user memory stats: {e}")
|
|
806
|
-
|
|
929
|
+
raise e
|
|
807
930
|
|
|
808
931
|
def upsert_user_memory(
|
|
809
932
|
self, memory: UserMemory, deserialize: Optional[bool] = True
|
|
@@ -849,7 +972,43 @@ class FirestoreDb(BaseDb):
|
|
|
849
972
|
|
|
850
973
|
except Exception as e:
|
|
851
974
|
log_error(f"Exception upserting user memory: {e}")
|
|
852
|
-
|
|
975
|
+
raise e
|
|
976
|
+
|
|
977
|
+
def upsert_memories(
|
|
978
|
+
self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
|
|
979
|
+
) -> List[Union[UserMemory, Dict[str, Any]]]:
|
|
980
|
+
"""
|
|
981
|
+
Bulk upsert multiple user memories for improved performance on large datasets.
|
|
982
|
+
|
|
983
|
+
Args:
|
|
984
|
+
memories (List[UserMemory]): List of memories to upsert.
|
|
985
|
+
deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
|
|
986
|
+
|
|
987
|
+
Returns:
|
|
988
|
+
List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
|
|
989
|
+
|
|
990
|
+
Raises:
|
|
991
|
+
Exception: If an error occurs during bulk upsert.
|
|
992
|
+
"""
|
|
993
|
+
if not memories:
|
|
994
|
+
return []
|
|
995
|
+
|
|
996
|
+
try:
|
|
997
|
+
log_info(
|
|
998
|
+
f"FirestoreDb doesn't support efficient bulk operations, falling back to individual upserts for {len(memories)} memories"
|
|
999
|
+
)
|
|
1000
|
+
# Fall back to individual upserts
|
|
1001
|
+
results = []
|
|
1002
|
+
for memory in memories:
|
|
1003
|
+
if memory is not None:
|
|
1004
|
+
result = self.upsert_user_memory(memory, deserialize=deserialize)
|
|
1005
|
+
if result is not None:
|
|
1006
|
+
results.append(result)
|
|
1007
|
+
return results
|
|
1008
|
+
|
|
1009
|
+
except Exception as e:
|
|
1010
|
+
log_error(f"Exception during bulk memory upsert: {e}")
|
|
1011
|
+
return []
|
|
853
1012
|
|
|
854
1013
|
def clear_memories(self) -> None:
|
|
855
1014
|
"""Delete all memories from the database.
|
|
@@ -882,9 +1041,215 @@ class FirestoreDb(BaseDb):
|
|
|
882
1041
|
batch.commit()
|
|
883
1042
|
|
|
884
1043
|
except Exception as e:
|
|
885
|
-
|
|
1044
|
+
log_error(f"Exception deleting all memories: {e}")
|
|
1045
|
+
raise e
|
|
1046
|
+
|
|
1047
|
+
# -- Cultural Knowledge methods --
|
|
1048
|
+
def clear_cultural_knowledge(self) -> None:
|
|
1049
|
+
"""Delete all cultural knowledge from the database.
|
|
1050
|
+
|
|
1051
|
+
Raises:
|
|
1052
|
+
Exception: If an error occurs during deletion.
|
|
1053
|
+
"""
|
|
1054
|
+
try:
|
|
1055
|
+
collection_ref = self._get_collection(table_type="culture")
|
|
1056
|
+
|
|
1057
|
+
# Get all documents in the collection
|
|
1058
|
+
docs = collection_ref.stream()
|
|
1059
|
+
|
|
1060
|
+
# Delete all documents in batches
|
|
1061
|
+
batch = self.db_client.batch()
|
|
1062
|
+
batch_count = 0
|
|
1063
|
+
|
|
1064
|
+
for doc in docs:
|
|
1065
|
+
batch.delete(doc.reference)
|
|
1066
|
+
batch_count += 1
|
|
1067
|
+
|
|
1068
|
+
# Firestore batch has a limit of 500 operations
|
|
1069
|
+
if batch_count >= 500:
|
|
1070
|
+
batch.commit()
|
|
1071
|
+
batch = self.db_client.batch()
|
|
1072
|
+
batch_count = 0
|
|
1073
|
+
|
|
1074
|
+
# Commit remaining operations
|
|
1075
|
+
if batch_count > 0:
|
|
1076
|
+
batch.commit()
|
|
1077
|
+
|
|
1078
|
+
except Exception as e:
|
|
1079
|
+
log_error(f"Exception deleting all cultural knowledge: {e}")
|
|
1080
|
+
raise e
|
|
1081
|
+
|
|
1082
|
+
def delete_cultural_knowledge(self, id: str) -> None:
|
|
1083
|
+
"""Delete cultural knowledge by ID.
|
|
1084
|
+
|
|
1085
|
+
Args:
|
|
1086
|
+
id (str): The ID of the cultural knowledge to delete.
|
|
1087
|
+
|
|
1088
|
+
Raises:
|
|
1089
|
+
Exception: If an error occurs during deletion.
|
|
1090
|
+
"""
|
|
1091
|
+
try:
|
|
1092
|
+
collection_ref = self._get_collection(table_type="culture")
|
|
1093
|
+
docs = collection_ref.where(filter=FieldFilter("id", "==", id)).stream()
|
|
1094
|
+
|
|
1095
|
+
for doc in docs:
|
|
1096
|
+
doc.reference.delete()
|
|
1097
|
+
log_debug(f"Deleted cultural knowledge with ID: {id}")
|
|
1098
|
+
|
|
1099
|
+
except Exception as e:
|
|
1100
|
+
log_error(f"Error deleting cultural knowledge: {e}")
|
|
1101
|
+
raise e
|
|
1102
|
+
|
|
1103
|
+
def get_cultural_knowledge(
|
|
1104
|
+
self, id: str, deserialize: Optional[bool] = True
|
|
1105
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
1106
|
+
"""Get cultural knowledge by ID.
|
|
1107
|
+
|
|
1108
|
+
Args:
|
|
1109
|
+
id (str): The ID of the cultural knowledge to retrieve.
|
|
1110
|
+
deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge object. Defaults to True.
|
|
1111
|
+
|
|
1112
|
+
Returns:
|
|
1113
|
+
Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The cultural knowledge if found, None otherwise.
|
|
1114
|
+
|
|
1115
|
+
Raises:
|
|
1116
|
+
Exception: If an error occurs during retrieval.
|
|
1117
|
+
"""
|
|
1118
|
+
try:
|
|
1119
|
+
collection_ref = self._get_collection(table_type="culture")
|
|
1120
|
+
docs = collection_ref.where(filter=FieldFilter("id", "==", id)).limit(1).stream()
|
|
1121
|
+
|
|
1122
|
+
for doc in docs:
|
|
1123
|
+
result = doc.to_dict()
|
|
1124
|
+
if not deserialize:
|
|
1125
|
+
return result
|
|
1126
|
+
return deserialize_cultural_knowledge_from_db(result)
|
|
1127
|
+
|
|
1128
|
+
return None
|
|
1129
|
+
|
|
1130
|
+
except Exception as e:
|
|
1131
|
+
log_error(f"Error getting cultural knowledge: {e}")
|
|
1132
|
+
raise e
|
|
1133
|
+
|
|
1134
|
+
def get_all_cultural_knowledge(
|
|
1135
|
+
self,
|
|
1136
|
+
agent_id: Optional[str] = None,
|
|
1137
|
+
team_id: Optional[str] = None,
|
|
1138
|
+
name: Optional[str] = None,
|
|
1139
|
+
limit: Optional[int] = None,
|
|
1140
|
+
page: Optional[int] = None,
|
|
1141
|
+
sort_by: Optional[str] = None,
|
|
1142
|
+
sort_order: Optional[str] = None,
|
|
1143
|
+
deserialize: Optional[bool] = True,
|
|
1144
|
+
) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
|
|
1145
|
+
"""Get all cultural knowledge with filtering and pagination.
|
|
1146
|
+
|
|
1147
|
+
Args:
|
|
1148
|
+
agent_id (Optional[str]): Filter by agent ID.
|
|
1149
|
+
team_id (Optional[str]): Filter by team ID.
|
|
1150
|
+
name (Optional[str]): Filter by name (case-insensitive partial match).
|
|
1151
|
+
limit (Optional[int]): Maximum number of results to return.
|
|
1152
|
+
page (Optional[int]): Page number for pagination.
|
|
1153
|
+
sort_by (Optional[str]): Field to sort by.
|
|
1154
|
+
sort_order (Optional[str]): Sort order ('asc' or 'desc').
|
|
1155
|
+
deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge objects. Defaults to True.
|
|
1156
|
+
|
|
1157
|
+
Returns:
|
|
1158
|
+
Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
|
|
1159
|
+
- When deserialize=True: List of CulturalKnowledge objects
|
|
1160
|
+
- When deserialize=False: Tuple with list of dictionaries and total count
|
|
1161
|
+
|
|
1162
|
+
Raises:
|
|
1163
|
+
Exception: If an error occurs during retrieval.
|
|
1164
|
+
"""
|
|
1165
|
+
try:
|
|
1166
|
+
collection_ref = self._get_collection(table_type="culture")
|
|
1167
|
+
|
|
1168
|
+
# Build query with filters
|
|
1169
|
+
query = collection_ref
|
|
1170
|
+
if agent_id is not None:
|
|
1171
|
+
query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
|
|
1172
|
+
if team_id is not None:
|
|
1173
|
+
query = query.where(filter=FieldFilter("team_id", "==", team_id))
|
|
1174
|
+
|
|
1175
|
+
# Get all matching documents
|
|
1176
|
+
docs = query.stream()
|
|
1177
|
+
results = [doc.to_dict() for doc in docs]
|
|
1178
|
+
|
|
1179
|
+
# Apply name filter (Firestore doesn't support regex in queries)
|
|
1180
|
+
if name is not None:
|
|
1181
|
+
results = [r for r in results if name.lower() in r.get("name", "").lower()]
|
|
1182
|
+
|
|
1183
|
+
total_count = len(results)
|
|
1184
|
+
|
|
1185
|
+
# Apply sorting and pagination to in-memory results
|
|
1186
|
+
sorted_results = apply_sorting_to_records(records=results, sort_by=sort_by, sort_order=sort_order)
|
|
1187
|
+
paginated_results = apply_pagination_to_records(records=sorted_results, limit=limit, page=page)
|
|
1188
|
+
|
|
1189
|
+
if not deserialize:
|
|
1190
|
+
return paginated_results, total_count
|
|
1191
|
+
|
|
1192
|
+
return [deserialize_cultural_knowledge_from_db(item) for item in paginated_results]
|
|
1193
|
+
|
|
1194
|
+
except Exception as e:
|
|
1195
|
+
log_error(f"Error getting all cultural knowledge: {e}")
|
|
1196
|
+
raise e
|
|
1197
|
+
|
|
1198
|
+
def upsert_cultural_knowledge(
|
|
1199
|
+
self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
|
|
1200
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
1201
|
+
"""Upsert cultural knowledge in Firestore.
|
|
1202
|
+
|
|
1203
|
+
Args:
|
|
1204
|
+
cultural_knowledge (CulturalKnowledge): The cultural knowledge to upsert.
|
|
1205
|
+
deserialize (Optional[bool]): Whether to deserialize the result. Defaults to True.
|
|
886
1206
|
|
|
887
|
-
|
|
1207
|
+
Returns:
|
|
1208
|
+
Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The upserted cultural knowledge.
|
|
1209
|
+
|
|
1210
|
+
Raises:
|
|
1211
|
+
Exception: If an error occurs during upsert.
|
|
1212
|
+
"""
|
|
1213
|
+
try:
|
|
1214
|
+
collection_ref = self._get_collection(table_type="culture", create_collection_if_not_found=True)
|
|
1215
|
+
|
|
1216
|
+
# Serialize content, categories, and notes into a dict for DB storage
|
|
1217
|
+
content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
|
|
1218
|
+
|
|
1219
|
+
# Create the update document with serialized content
|
|
1220
|
+
update_doc = {
|
|
1221
|
+
"id": cultural_knowledge.id,
|
|
1222
|
+
"name": cultural_knowledge.name,
|
|
1223
|
+
"summary": cultural_knowledge.summary,
|
|
1224
|
+
"content": content_dict if content_dict else None,
|
|
1225
|
+
"metadata": cultural_knowledge.metadata,
|
|
1226
|
+
"input": cultural_knowledge.input,
|
|
1227
|
+
"created_at": cultural_knowledge.created_at,
|
|
1228
|
+
"updated_at": int(time.time()),
|
|
1229
|
+
"agent_id": cultural_knowledge.agent_id,
|
|
1230
|
+
"team_id": cultural_knowledge.team_id,
|
|
1231
|
+
}
|
|
1232
|
+
|
|
1233
|
+
# Find and update or create new document
|
|
1234
|
+
docs = collection_ref.where(filter=FieldFilter("id", "==", cultural_knowledge.id)).limit(1).stream()
|
|
1235
|
+
|
|
1236
|
+
doc_found = False
|
|
1237
|
+
for doc in docs:
|
|
1238
|
+
doc.reference.set(update_doc)
|
|
1239
|
+
doc_found = True
|
|
1240
|
+
break
|
|
1241
|
+
|
|
1242
|
+
if not doc_found:
|
|
1243
|
+
collection_ref.add(update_doc)
|
|
1244
|
+
|
|
1245
|
+
if not deserialize:
|
|
1246
|
+
return update_doc
|
|
1247
|
+
|
|
1248
|
+
return deserialize_cultural_knowledge_from_db(update_doc)
|
|
1249
|
+
|
|
1250
|
+
except Exception as e:
|
|
1251
|
+
log_error(f"Error upserting cultural knowledge: {e}")
|
|
1252
|
+
raise e
|
|
888
1253
|
|
|
889
1254
|
# -- Metrics methods --
|
|
890
1255
|
|
|
@@ -918,8 +1283,8 @@ class FirestoreDb(BaseDb):
|
|
|
918
1283
|
return results
|
|
919
1284
|
|
|
920
1285
|
except Exception as e:
|
|
921
|
-
log_error(f"Exception
|
|
922
|
-
|
|
1286
|
+
log_error(f"Exception getting all sessions for metrics calculation: {e}")
|
|
1287
|
+
raise e
|
|
923
1288
|
|
|
924
1289
|
def _get_metrics_calculation_starting_date(self, collection_ref) -> Optional[date]:
|
|
925
1290
|
"""Get the first date for which metrics calculation is needed."""
|
|
@@ -951,7 +1316,7 @@ class FirestoreDb(BaseDb):
|
|
|
951
1316
|
|
|
952
1317
|
except Exception as e:
|
|
953
1318
|
log_error(f"Exception getting metrics calculation starting date: {e}")
|
|
954
|
-
|
|
1319
|
+
raise e
|
|
955
1320
|
|
|
956
1321
|
def calculate_metrics(self) -> Optional[list[dict]]:
|
|
957
1322
|
"""Calculate metrics for all dates without complete metrics."""
|
|
@@ -1043,7 +1408,7 @@ class FirestoreDb(BaseDb):
|
|
|
1043
1408
|
|
|
1044
1409
|
except Exception as e:
|
|
1045
1410
|
log_error(f"Exception getting metrics: {e}")
|
|
1046
|
-
|
|
1411
|
+
raise e
|
|
1047
1412
|
|
|
1048
1413
|
# -- Knowledge methods --
|
|
1049
1414
|
|
|
@@ -1064,7 +1429,8 @@ class FirestoreDb(BaseDb):
|
|
|
1064
1429
|
doc.reference.delete()
|
|
1065
1430
|
|
|
1066
1431
|
except Exception as e:
|
|
1067
|
-
log_error(f"Error deleting knowledge
|
|
1432
|
+
log_error(f"Error deleting knowledge content: {e}")
|
|
1433
|
+
raise e
|
|
1068
1434
|
|
|
1069
1435
|
def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
|
|
1070
1436
|
"""Get a knowledge row from the database.
|
|
@@ -1089,8 +1455,8 @@ class FirestoreDb(BaseDb):
|
|
|
1089
1455
|
return None
|
|
1090
1456
|
|
|
1091
1457
|
except Exception as e:
|
|
1092
|
-
log_error(f"Error getting knowledge
|
|
1093
|
-
|
|
1458
|
+
log_error(f"Error getting knowledge content: {e}")
|
|
1459
|
+
raise e
|
|
1094
1460
|
|
|
1095
1461
|
def get_knowledge_contents(
|
|
1096
1462
|
self,
|
|
@@ -1138,8 +1504,8 @@ class FirestoreDb(BaseDb):
|
|
|
1138
1504
|
return knowledge_rows, total_count
|
|
1139
1505
|
|
|
1140
1506
|
except Exception as e:
|
|
1141
|
-
log_error(f"Error getting knowledge
|
|
1142
|
-
|
|
1507
|
+
log_error(f"Error getting knowledge contents: {e}")
|
|
1508
|
+
raise e
|
|
1143
1509
|
|
|
1144
1510
|
def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
|
|
1145
1511
|
"""Upsert knowledge content in the database.
|
|
@@ -1169,8 +1535,8 @@ class FirestoreDb(BaseDb):
|
|
|
1169
1535
|
return knowledge_row
|
|
1170
1536
|
|
|
1171
1537
|
except Exception as e:
|
|
1172
|
-
log_error(f"Error upserting knowledge
|
|
1173
|
-
|
|
1538
|
+
log_error(f"Error upserting knowledge content: {e}")
|
|
1539
|
+
raise e
|
|
1174
1540
|
|
|
1175
1541
|
# -- Eval methods --
|
|
1176
1542
|
|
|
@@ -1193,7 +1559,7 @@ class FirestoreDb(BaseDb):
|
|
|
1193
1559
|
|
|
1194
1560
|
except Exception as e:
|
|
1195
1561
|
log_error(f"Error creating eval run: {e}")
|
|
1196
|
-
|
|
1562
|
+
raise e
|
|
1197
1563
|
|
|
1198
1564
|
def delete_eval_run(self, eval_run_id: str) -> None:
|
|
1199
1565
|
"""Delete an eval run from the database."""
|
|
@@ -1213,7 +1579,7 @@ class FirestoreDb(BaseDb):
|
|
|
1213
1579
|
|
|
1214
1580
|
except Exception as e:
|
|
1215
1581
|
log_error(f"Error deleting eval run {eval_run_id}: {e}")
|
|
1216
|
-
raise
|
|
1582
|
+
raise e
|
|
1217
1583
|
|
|
1218
1584
|
def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
|
|
1219
1585
|
"""Delete multiple eval runs from the database.
|
|
@@ -1244,7 +1610,7 @@ class FirestoreDb(BaseDb):
|
|
|
1244
1610
|
|
|
1245
1611
|
except Exception as e:
|
|
1246
1612
|
log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
|
|
1247
|
-
raise
|
|
1613
|
+
raise e
|
|
1248
1614
|
|
|
1249
1615
|
def get_eval_run(
|
|
1250
1616
|
self, eval_run_id: str, deserialize: Optional[bool] = True
|
|
@@ -1265,6 +1631,9 @@ class FirestoreDb(BaseDb):
|
|
|
1265
1631
|
"""
|
|
1266
1632
|
try:
|
|
1267
1633
|
collection_ref = self._get_collection(table_type="evals")
|
|
1634
|
+
if not collection_ref:
|
|
1635
|
+
return None
|
|
1636
|
+
|
|
1268
1637
|
docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
|
|
1269
1638
|
|
|
1270
1639
|
eval_run_raw = None
|
|
@@ -1282,7 +1651,7 @@ class FirestoreDb(BaseDb):
|
|
|
1282
1651
|
|
|
1283
1652
|
except Exception as e:
|
|
1284
1653
|
log_error(f"Exception getting eval run {eval_run_id}: {e}")
|
|
1285
|
-
|
|
1654
|
+
raise e
|
|
1286
1655
|
|
|
1287
1656
|
def get_eval_runs(
|
|
1288
1657
|
self,
|
|
@@ -1383,7 +1752,7 @@ class FirestoreDb(BaseDb):
|
|
|
1383
1752
|
|
|
1384
1753
|
except Exception as e:
|
|
1385
1754
|
log_error(f"Exception getting eval runs: {e}")
|
|
1386
|
-
|
|
1755
|
+
raise e
|
|
1387
1756
|
|
|
1388
1757
|
def rename_eval_run(
|
|
1389
1758
|
self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
|
|
@@ -1405,6 +1774,8 @@ class FirestoreDb(BaseDb):
|
|
|
1405
1774
|
"""
|
|
1406
1775
|
try:
|
|
1407
1776
|
collection_ref = self._get_collection(table_type="evals")
|
|
1777
|
+
if not collection_ref:
|
|
1778
|
+
return None
|
|
1408
1779
|
|
|
1409
1780
|
docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
|
|
1410
1781
|
doc_ref = next((doc.reference for doc in docs), None)
|
|
@@ -1429,4 +1800,4 @@ class FirestoreDb(BaseDb):
|
|
|
1429
1800
|
|
|
1430
1801
|
except Exception as e:
|
|
1431
1802
|
log_error(f"Error updating eval run name {eval_run_id}: {e}")
|
|
1432
|
-
raise
|
|
1803
|
+
raise e
|