from urllib import request
from alayaflow.clients.alayamem.http_client import HttpAlayaMemClient


class RetrieveComponent:
    def __init__(self, client: HttpAlayaMemClient):
        self.client = client
    
    def __call__(self, query: str, collection_name: str, limit: int = 3) -> list[str]:
        result = self.client.vdb_query([query], limit, collection_name)
        return result.get('documents', [[]])[0] if result.get('documents') else []

if __name__ == "__main__":
    client = HttpAlayaMemClient("http://10.16.70.46:5555")
    res = client.vdb_query(messages="姓名", limit=5, collection_name="file_watcher_collection")

    print(res)