app/services/embedding_service.py

from app.repositories.neo4j_repository import (
    Neo4jRepository,
)

from app.services.embedding_model_service import (
    EmbeddingModelService,
)


class EmbeddingService:

    def __init__(self):

        self.embedding_model = (
            EmbeddingModelService()
        )

        self.neo = Neo4jRepository()

    def build(
        self,
        repository_id: str,
        resolved,
    ):

        for file in resolved:

            self._embed_functions(
                repository_id,
                file["functions"],
            )

            self._embed_classes(
                repository_id,
                file["classes"],
            )

    def _embed_functions(
        self,
        repository_id: str,
        functions,
    ):

        for function in functions:

            source = function.get(
                "source_code"
            )

            if not source:
                continue

            embedding = (
                self.embedding_model.encode(
                    source
                )
            )

            self.neo.update_embedding(
                repository_id=repository_id,
                qualified_name=function[
                    "qualified_name"
                ],
                embedding=embedding,
            )

    def _embed_classes(
        self,
        repository_id: str,
        classes,
    ):

        for cls in classes:

            source = cls.get(
                "source_code"
            )

            if source:

                embedding = (
                    self.embedding_model.encode(
                        source
                    )
                )

                self.neo.update_embedding(
                    repository_id=repository_id,
                    qualified_name=cls[
                        "qualified_name"
                    ],
                    embedding=embedding,
                )

            self._embed_functions(
                repository_id,
                cls["methods"],
            )