Repository: StanGirard/quivr
File: backend/api/quivr_api/modules/knowledge/service/knowledge_service.py (Python)

Code before:
from typing import List
from uuid import UUID

from quivr_core.models import QuivrKnowledge as Knowledge

from quivr_api.logger import get_logger
from quivr_api.modules.dependencies import BaseService
from quivr_api.modules.knowledge.dto.inputs import (
    CreateKnowledgeProperties,
    KnowledgeStatus,
)
from quivr_api.modules.knowledge.dto.outputs import DeleteKnowledgeResponse
from quivr_api.modules.knowledge.entity.knowledge import KnowledgeDB
from quivr_api.modules.knowledge.entity.knowledge_brain import KnowledgeBrain
from quivr_api.modules.knowledge.repository.knowledges import KnowledgeRepository

logger = get_logger(__name__)


class KnowledgeService(BaseService[KnowledgeRepository]):
    repository_cls = KnowledgeRepository

    def __init__(self, repository: KnowledgeRepository):
        self.repository = repository

    async def add_knowledge(
        self,
        knowledge_to_add: CreateKnowledgeProperties,  # FIXME: (later) @Amine brain id should not be in CreateKnowledgeProperties but since storage is brain_id/file_name
    ) -> Knowledge:
        knowledge = KnowledgeDB(
            file_name=knowledge_to_add.file_name,
            url=knowledge_to_add.url,
            mime_type=knowledge_to_add.mime_type,
            status=knowledge_to_add.status.value,
            source=knowledge_to_add.source,
            source_link=knowledge_to_add.source_link,
            file_size=knowledge_to_add.file_size,
            file_sha1=knowledge_to_add.file_sha1,
            metadata_=knowledge_to_add.metadata,  # type: ignore
        )

        inserted_knowledge_db_instance = await self.repository.insert_knowledge(
            knowledge
        )

        assert inserted_knowledge_db_instance.id, "Knowledge ID not generated"
        if inserted_knowledge_db_instance.source == "local":
            source_link = f"s3://quivr/{knowledge_to_add.brain_id}/{inserted_knowledge_db_instance.id}"

        inserted_knowledge = await self.repository.update_source_link_knowledge(
            knowledge_id=inserted_knowledge_db_instance.id, source_link=source_link

Review comment (suggestion):
We should be careful here. If we return a KnowledgeDB object you should refresh the object using `await self.session.refresh(obj)` to have the inserted object fields

Code after:
from typing import List
from uuid import UUID

from quivr_core.models import KnowledgeStatus

from quivr_api.logger import get_logger
from quivr_api.modules.dependencies import BaseService
from quivr_api.modules.knowledge.dto.inputs import (
    CreateKnowledgeProperties,
)
from quivr_api.modules.knowledge.dto.outputs import DeleteKnowledgeResponse
from quivr_api.modules.knowledge.entity.knowledge import Knowledge, KnowledgeDB
from quivr_api.modules.knowledge.repository.knowledges import KnowledgeRepository

logger = get_logger(__name__)


class KnowledgeService(BaseService[KnowledgeRepository]):
    repository_cls = KnowledgeRepository

    def __init__(self, repository: KnowledgeRepository):
        self.repository = repository

    async def add_knowledge(
        self,
        knowledge_to_add: CreateKnowledgeProperties,  # FIXME: (later) @Amine brain id should not be in CreateKnowledgeProperties but since storage is brain_id/file_name
    ) -> Knowledge:
        knowledge = KnowledgeDB(
            file_name=knowledge_to_add.file_name,
            url=knowledge_to_add.url,
            mime_type=knowledge_to_add.mime_type,
            status=knowledge_to_add.status.value,
            source=knowledge_to_add.source,
            source_link=knowledge_to_add.source_link,
            file_size=knowledge_to_add.file_size,
            file_sha1=knowledge_to_add.file_sha1,
            metadata_=knowledge_to_add.metadata,  # type: ignore
        )

        inserted_knowledge_db_instance = await self.repository.insert_knowledge(
            knowledge, brain_id=knowledge_to_add.brain_id
        )

        assert inserted_knowledge_db_instance.id, "Knowledge ID not generated"
        if inserted_knowledge_db_instance.source == "local":
            source_link = f"s3://quivr/{knowledge_to_add.brain_id}/{inserted_knowledge_db_instance.id}"

        inserted_knowledge = await self.repository.update_source_link_knowledge(
            knowledge_id=inserted_knowledge_db_instance.id, source_link=source_link
        )