forked from innovacion/searchbox
68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
import subprocess
|
|
import time
|
|
|
|
import pytest
|
|
from fastembed import TextEmbedding
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, PointStruct, VectorParams
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def embedding_model():
|
|
return TextEmbedding()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def qdrant_client(embedding_model: TextEmbedding):
|
|
client = QdrantClient(":memory:")
|
|
|
|
documents: list[str] = [
|
|
"Rick es el mas guapo",
|
|
"Los pulpos tienen tres corazones y sangre azul",
|
|
"Las cucarachas pueden vivir hasta una semana sin cabeza",
|
|
"Los koalas tienen huellas dactilares casi idénticas a las humanas",
|
|
"La miel nunca se echa a perder, incluso después de miles de años",
|
|
]
|
|
embeddings = list(embedding_model.embed(documents))
|
|
size = len(embeddings[0])
|
|
|
|
_ = client.recreate_collection(
|
|
collection_name="dummy_collection",
|
|
vectors_config=VectorParams(distance=Distance.COSINE, size=size),
|
|
)
|
|
|
|
for idx, (emb, document) in enumerate(zip(embeddings, documents)):
|
|
_ = client.upsert(
|
|
collection_name="dummy_collection",
|
|
points=[
|
|
PointStruct(id=idx, vector=emb.tolist(), payload={"text": document})
|
|
],
|
|
)
|
|
|
|
yield client
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def run_mcp():
|
|
# Start the MCP server in the background
|
|
process = subprocess.Popen(
|
|
["uv", "run", "vector-search-mcp"],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
|
|
# Give the server a moment to start up
|
|
time.sleep(2)
|
|
|
|
try:
|
|
yield "http://localhost:8000/sse"
|
|
finally:
|
|
# Clean up the process when tests are done
|
|
process.terminate()
|
|
try:
|
|
_ = process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
_ = process.wait()
|