70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import asyncio
|
|
import time
|
|
|
|
import aiohttp
|
|
import requests
|
|
|
|
from app.conf.app_config import EmbeddingConfig, app_config
|
|
|
|
|
|
class EmbeddingClientManager:
|
|
def __init__(self, config: EmbeddingConfig):
|
|
self.config = config
|
|
self.client = None
|
|
|
|
def _get_url(self):
|
|
return f"http://{self.config.host}:{self.config.port}"
|
|
|
|
|
|
def init(self, wait_for_ready: bool = True):
|
|
self.client = aiohttp.ClientSession()
|
|
|
|
async def close(self):
|
|
if self.client:
|
|
await self.client.close()
|
|
self.client = None
|
|
|
|
async def aembed_documents(self, texts: list) -> list:
|
|
if not self.client:
|
|
self.init(wait_for_ready=False)
|
|
|
|
url = f"{self._get_url()}/embed"
|
|
payload = {
|
|
"inputs": texts,
|
|
"parameters": {"truncate": True}
|
|
}
|
|
|
|
async with self.client.post(url, json=payload) as response:
|
|
result = await response.json()
|
|
if isinstance(result, list):
|
|
return result
|
|
return result.get("embeddings", [])
|
|
|
|
async def aembed_query(self, text: str) -> list:
|
|
embeddings = await self.aembed_documents([text])
|
|
return embeddings[0] if embeddings else []
|
|
|
|
|
|
embedding_client = EmbeddingClientManager(app_config.embedding)
|
|
|
|
if __name__ == "__main__":
|
|
print("Testing EmbeddingClientManager...")
|
|
|
|
async def test():
|
|
try:
|
|
embedding_client.init(wait_for_ready=True)
|
|
print("Initialization successful")
|
|
|
|
text = "What is Deep Learning?"
|
|
print(f"Test text: {text}")
|
|
|
|
result = await embedding_client.aembed_query(text)
|
|
print(f"Embedding successful")
|
|
print(f"Vector length: {len(result)}")
|
|
print(f"First 10 values: {result[:10]}")
|
|
|
|
await embedding_client.close()
|
|
except Exception as e:
|
|
print(f"Test failed: {str(e)}")
|
|
|
|
asyncio.run(test()) |