ADK python , RAG Search

.env

GOOGLE_GENAI_USE_VERTEXAI=FALSE
GOOGLE_API_KEY=your-key

__init__.py

from . import agent

agent.py

# rag_search/agent.py

import datetime
import ollama
import os
import psycopg2
import json # ベクトルを文字列として渡すために使用

from dotenv import load_dotenv
from google.adk.agents import Agent
from pgvector.psycopg2 import register_vector
from zoneinfo import ZoneInfo

#
def search_similar(embedding, top_k=5):
    # PostgreSQL 接続
    conn = psycopg2.connect(
        dbname="mydb",
        user="root",
        password="admin",
        host="localhost",
        port=5432
    )

    # pgvector を psycopg2 に登録
    register_vector(conn)

    # -- コサイン距離をコサイン類似度 (1 - 距離) に変換し、similarityとして取得
    # -- 類似度が高い順（降順）に並び替え
    SQL_QUERY = """
    SELECT
        id,
        content,
        embedding,
        1 - (embedding <=> %s) AS cosine_similarity
    FROM
        documents
    ORDER BY
        cosine_similarity DESC
    LIMIT 3;
    """

    query_vector_str = json.dumps(embedding)
    with conn.cursor() as cur:
        cur.execute(SQL_QUERY, (query_vector_str,))
        return cur.fetchall()

#
def rag_search(query: str) -> dict:
    """入力文字から、RAG検索します。

    Args:
        query (str): RAG検索のクエリ文字列。

    Returns:
        dict: RAG検索の結果
    """
    #vec
    # --- 環境変数の取得 ---
    embedding = ollama.embeddings(
        model="qwen3-embedding:0.6b",
        prompt=query
    )

    vec = embedding["embedding"]

    print(len(vec))
    print(vec[:5])    

    matches = ""
    rows = search_similar(vec, 3)
    ouStr = "" 

    for row in rows:
        ouStr += row[1]+ "\n\n" 

    #print(ouStr)

    return {
        "status": "success",
        "report": f" {ouStr}"
    }


root_agent = Agent(
    name="weather_time_agent",
    model="gemini-2.0-flash",
    description=(
        "RAG検索クエリーに回答に答えるエージェント。"
    ),
    instruction=(
        "あなたは、ユーザーの検索依頼に答えることができる親切なエージェントです。"
    ),
    tools=[rag_search],
)