retrieve.py

import ollama
import os
from supabase import create_client
from dotenv import load_dotenv

def embed_query(user_query):
    response = ollama.embed(model='qwen3-embedding:4b', input=user_query)
    return response['embeddings'][0]

def match_chunks(supabase, embedding, top_k=3):
    response = supabase.rpc('match_chunks', {'query_embedding' :embedding, 'match_count' :top_k}).execute()
    return response
    # 返回对象：
    # APIResponse(
#     data=[
#         {'content': 'xxx', 'similarity': 0.85},
#         {'content': 'yyy', 'similarity': 0.82},
#         ...
#     ],
#     count=None
#     )

def main():
    load_dotenv() # 把.env 文件里的内容加载到环境变量里

    SUPABASE_URL = os.getenv('SUPABASE_URL') # 从环境变量里读取
    SUPABASE_KEY = os.getenv('SUPABASE_KEY')
    supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

    user_query = 'test_query'
    user_embed = embed_query(user_query)
    response = match_chunks(supabase, user_embed)
    return response

if __name__ == '__main__':
    main()