from app.repositories.neo4j_repository import (
    Neo4jRepository,
)

from app.services.embedding_model_service import (
    EmbeddingModelService,
)


class RetrievalService:

    def __init__(self):

        self.embedding = EmbeddingModelService()

        self.neo = Neo4jRepository()

    def search(
        self,
        repository_id: str,
        query: str,
        limit: int = 5,
        hops: int = 1,
    ):

        vector = self.embedding.encode(
            query
        )

        return self.neo.retrieve_context(
            repository_id=repository_id,
            embedding=vector,
            limit=limit,
            hops=hops,
        )