chore: Fixes test_fast_api.py (part I for circular deps).

It still fails due to signal used not in main thread. It will be fixed later.

PiperOrigin-RevId: 760050504
This commit is contained in:
Wei Sun (Jack)
2025-05-17 12:21:22 -07:00
committed by Copybara-Service
parent f592de4cc0
commit 9324801b75
2 changed files with 64 additions and 39 deletions

View File

@@ -12,23 +12,29 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from collections import OrderedDict
import json
import os
import tempfile
from typing import Optional
from typing import TYPE_CHECKING
from google.genai import types
from typing_extensions import override
from vertexai.preview import rag
from ..events.event import Event
from ..sessions.session import Session
from . import _utils
from .base_memory_service import BaseMemoryService
from .base_memory_service import SearchMemoryResponse
from .memory_entry import MemoryEntry
if TYPE_CHECKING:
from ..events.event import Event
from ..sessions.session import Session
class VertexAiRagMemoryService(BaseMemoryService):
"""A memory service that uses Vertex AI RAG for storage and retrieval."""
@@ -103,6 +109,8 @@ class VertexAiRagMemoryService(BaseMemoryService):
self, *, app_name: str, user_id: str, query: str
) -> SearchMemoryResponse:
"""Searches for sessions that match the query using rag.retrieval_query."""
from ..events.event import Event
response = rag.retrieval_query(
text=query,
rag_resources=self._vertex_rag_store.rag_resources,