fix: Knowledge

Pipfile

@@ -0,0 +1,11 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+
+[dev-packages]
+
+[requires]
+python_version = "3.11"

backend/api/quivr_api/modules/brain/integrations/Multi_Contract/Brain.py

@@ -2,7 +2,8 @@
 from operator import itemgetter
 from typing import List
 
-from langchain.prompts import HumanMessagePromptTemplate, SystemMessagePromptTemplate
+from langchain.prompts import (HumanMessagePromptTemplate,
+                               SystemMessagePromptTemplate)
 from langchain_community.chat_models import ChatLiteLLM
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
@@ -115,7 +116,7 @@ def __init__(
     def get_chain(self):
 
         list_files_array = (
-            self.knowledge_qa.knowledge_service.get_all_knowledge_in_brain(
+            await self.knowledge_qa.knowledge_service.get_all_knowledge_in_brain(
                 self.brain_id
             )
         )  # pyright: ignore reportPrivateUsage=none

backend/api/quivr_api/modules/brain/service/brain_service.py

@@ -4,25 +4,23 @@
 from fastapi import HTTPException
 from quivr_api.celery_config import celery
 from quivr_api.logger import get_logger
-from quivr_api.modules.brain.dto.inputs import (
-    BrainUpdatableProperties,
-    CreateBrainProperties,
-)
+from quivr_api.modules.brain.dto.inputs import (BrainUpdatableProperties,
+                                                CreateBrainProperties)
 from quivr_api.modules.brain.entity.brain_entity import BrainEntity, BrainType
 from quivr_api.modules.brain.entity.integration_brain import IntegrationEntity
-from quivr_api.modules.brain.repository import (
-    Brains,
-    BrainsUsers,
-    BrainsVectors,
-    IntegrationBrain,
-    IntegrationDescription,
-)
-from quivr_api.modules.knowledge.service.knowledge_service import KnowledgeService
+from quivr_api.modules.brain.repository import (Brains, BrainsUsers,
+                                                BrainsVectors,
+                                                IntegrationBrain,
+                                                IntegrationDescription)
+from quivr_api.modules.dependencies import get_service
+from quivr_api.modules.knowledge.service.knowledge_service import \
+    KnowledgeService
 from quivr_api.vectorstore.supabase import CustomSupabaseVectorStore
 
 logger = get_logger(__name__)
 
-knowledge_service = KnowledgeService()
+#knowledge_service = KnowledgeService()
+knowledge_service = get_service(KnowledgeService)()
 
 
 class BrainService:

backend/api/quivr_api/modules/chat/controller/chat_routes.py

@@ -3,26 +3,23 @@
 
 from fastapi import APIRouter, Depends, HTTPException, Query, Request
 from fastapi.responses import StreamingResponse
-
 from quivr_api.logger import get_logger
 from quivr_api.middlewares.auth import AuthBearer, get_current_user
 from quivr_api.modules.brain.entity.brain_entity import RoleEnum
-from quivr_api.modules.brain.service.brain_authorization_service import (
-    validate_brain_authorization,
-)
+from quivr_api.modules.brain.service.brain_authorization_service import \
+    validate_brain_authorization
 from quivr_api.modules.brain.service.brain_service import BrainService
 from quivr_api.modules.chat.dto.chats import ChatItem, ChatQuestion
-from quivr_api.modules.chat.dto.inputs import (
-    ChatMessageProperties,
-    ChatUpdatableProperties,
-    CreateChatProperties,
-    QuestionAndAnswer,
-)
+from quivr_api.modules.chat.dto.inputs import (ChatMessageProperties,
+                                               ChatUpdatableProperties,
+                                               CreateChatProperties,
+                                               QuestionAndAnswer)
 from quivr_api.modules.chat.entity.chat import Chat
 from quivr_api.modules.chat.service.chat_service import ChatService
 from quivr_api.modules.chat_llm_service.chat_llm_service import ChatLLMService
 from quivr_api.modules.dependencies import get_service
-from quivr_api.modules.knowledge.repository.knowledges import KnowledgeRepository
+from quivr_api.modules.knowledge.service.knowledge_service import \
+    KnowledgeService
 from quivr_api.modules.models.service.model_service import ModelService
 from quivr_api.modules.prompt.service.prompt_service import PromptService
 from quivr_api.modules.rag_service import RAGService
@@ -34,7 +31,9 @@
 
 chat_router = APIRouter()
 brain_service = BrainService()
-knowledge_service = KnowledgeRepository()
+KnowledgeServiceDep = Annotated[
+    KnowledgeService, Depends(get_service(KnowledgeService))
+]
 prompt_service = PromptService()
 
 
@@ -166,6 +165,7 @@ async def create_question_handler(
     chat_id: UUID,
     current_user: UserIdentityDep,
     chat_service: ChatServiceDep,
+    knowledge_service: KnowledgeServiceDep,
     model_service: ModelServiceDep,
     brain_id: Annotated[UUID | None, Query()] = None,
 ):
@@ -231,6 +231,7 @@ async def create_stream_question_handler(
     chat_id: UUID,
     chat_service: ChatServiceDep,
     current_user: UserIdentityDep,
+    knowledge_service: KnowledgeServiceDep,
     model_service: ModelServiceDep,
     brain_id: Annotated[UUID | None, Query()] = None,
 ) -> StreamingResponse:

backend/api/quivr_api/modules/knowledge/controller/knowledge_routes.py

@@ -1,6 +1,8 @@
+from typing import Annotated
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, HTTPException, Query
+
 from quivr_api.logger import get_logger
 from quivr_api.middlewares.auth import AuthBearer, get_current_user
 from quivr_api.modules.brain.entity.brain_entity import RoleEnum
@@ -9,6 +11,7 @@
     validate_brain_authorization,
 )
 from quivr_api.modules.brain.service.brain_vector_service import BrainVectorService
+from quivr_api.modules.dependencies import get_service
 from quivr_api.modules.knowledge.service.knowledge_service import KnowledgeService
 from quivr_api.modules.upload.service.generate_file_signed_url import (
     generate_file_signed_url,
@@ -18,13 +21,16 @@
 knowledge_router = APIRouter()
 logger = get_logger(__name__)
 
-knowledge_service = KnowledgeService()
+KnowledgeServiceDep = Annotated[
+    KnowledgeService, Depends(get_service(KnowledgeService))
+]
 
 
 @knowledge_router.get(
     "/knowledge", dependencies=[Depends(AuthBearer())], tags=["Knowledge"]
 )
 async def list_knowledge_in_brain_endpoint(
+    knowledge_service: KnowledgeServiceDep,
     brain_id: UUID = Query(..., description="The ID of the brain"),
     current_user: UserIdentity = Depends(get_current_user),
 ):
@@ -33,8 +39,11 @@ async def list_knowledge_in_brain_endpoint(
     """
 
     validate_brain_authorization(brain_id=brain_id, user_id=current_user.id)
+    logger.debug(
+        f"REPO {knowledge_service.repository}, REPO SESSION {knowledge_service.repository.session}"
+    )
 
-    knowledges = knowledge_service.get_all_knowledge(brain_id)
+    knowledges = await knowledge_service.get_all_knowledge(brain_id)
 
     return {"knowledges": knowledges}
 
@@ -49,16 +58,17 @@ async def list_knowledge_in_brain_endpoint(
 )
 async def delete_endpoint(
     knowledge_id: UUID,
+    knowledge_service: KnowledgeServiceDep,
     current_user: UserIdentity = Depends(get_current_user),
     brain_id: UUID = Query(..., description="The ID of the brain"),
 ):
     """
     Delete a specific knowledge from a brain.
     """
 
-    knowledge = knowledge_service.get_knowledge(knowledge_id)
+    knowledge = await knowledge_service.get_knowledge(knowledge_id)
     file_name = knowledge.file_name if knowledge.file_name else knowledge.url
-    knowledge_service.remove_knowledge(knowledge_id)
+    await knowledge_service.remove_knowledge(knowledge_id)
 
     brain_vector_service = BrainVectorService(brain_id)
     if knowledge.file_name:
@@ -78,13 +88,14 @@ async def delete_endpoint(
 )
 async def generate_signed_url_endpoint(
     knowledge_id: UUID,
+    knowledge_service: KnowledgeServiceDep,
     current_user: UserIdentity = Depends(get_current_user),
 ):
     """
     Generate a signed url to download the file from storage.
     """
 
-    knowledge = knowledge_service.get_knowledge(knowledge_id)
+    knowledge = await knowledge_service.get_knowledge(knowledge_id)
 
     validate_brain_authorization(brain_id=knowledge.brain_id, user_id=current_user.id)
 

backend/api/quivr_api/modules/knowledge/dto/inputs.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Optional
+from typing import Dict, Optional
 from uuid import UUID
 
 from pydantic import BaseModel
@@ -9,16 +9,17 @@ class KnowledgeStatus(str, Enum):
     PROCESSING = "PROCESSING"
     UPLOADED = "UPLOADED"
     ERROR = "ERROR"
-
-
 class CreateKnowledgeProperties(BaseModel):
     brain_id: UUID
     file_name: Optional[str] = None
     url: Optional[str] = None
-    extension: str = "txt"
-    integration: Optional[str] = None
-    integration_link: Optional[str] = None
+    mime_type: str = "txt" 
     status: KnowledgeStatus = KnowledgeStatus.PROCESSING
+    source: Optional[str] = None
+    source_link: Optional[str] = None 
+    file_size: Optional[int] = None 
+    file_sha1: Optional[str] = None  
+    metadata: Optional[Dict[str, str]] = None
 
     def dict(self, *args, **kwargs):
         knowledge_dict = super().dict(*args, **kwargs)

backend/api/quivr_api/modules/knowledge/entity/knowledge.py

@@ -1,7 +1,11 @@
-from typing import Optional
+from datetime import datetime
+from typing import Dict, Optional
 from uuid import UUID
 
 from pydantic import BaseModel
+from sqlalchemy import JSON, TIMESTAMP, Column, text
+from sqlmodel import UUID as PGUUID
+from sqlmodel import Field, SQLModel
 
 
 class Knowledge(BaseModel):
@@ -11,5 +15,54 @@ class Knowledge(BaseModel):
     url: Optional[str] = None
     extension: str = "txt"
     status: str
-    integration: Optional[str] = None
-    integration_link: Optional[str] = None
+    source: Optional[str] = None
+    source_link: Optional[str] = None
+    file_size: Optional[int] = None
+    file_sha1: Optional[str] = None
+    updated_at: Optional[datetime] = None
+    created_at: Optional[datetime] = None
+    metadata: Optional[Dict[str, str]] = None
+
+    def dict(self, *args, **kwargs):
+        knowledge_dict = super().dict(*args, **kwargs)
+        knowledge_dict["brain_id"] = str(knowledge_dict.get("brain_id"))
+        return knowledge_dict
+
+
+class KnowledgeDB(SQLModel, table=True):
+    __tablename__ = "knowledge"  # type: ignore
+
+    id: UUID | None = Field(
+        default=None,
+        sa_column=Column(
+            PGUUID,
+            server_default=text("uuid_generate_v4()"),
+            primary_key=True,
+        ),
+    )
+    brain_id: UUID = Field(nullable=False)
+    file_name: Optional[str] = Field(default=None, max_length=255)
+    url: Optional[str] = Field(default=None, max_length=2048)
+    mime_type: str = Field(default="txt", max_length=100)
+    status: str = Field(max_length=50)
+    source: str = Field(max_length=255)
+    source_link: str = Field(max_length=2048, default = f"s3://quivr/{brain_id}/{id}")
+    file_size: Optional[int] = Field(gt=0)  # FIXME: Should not be optional @chloedia
+    file_sha1: Optional[str] = Field(
+        max_length=40
+    )  # FIXME: Should not be optional @chloedia
+    updated_at: datetime | None = Field(
+        default=None,
+        sa_column=Column(
+            TIMESTAMP(timezone=False),
+            server_default=text("CURRENT_TIMESTAMP"),
+        ),
+    )
+    created_at: datetime | None = Field(
+        default=None,
+        sa_column=Column(
+            TIMESTAMP(timezone=False),
+            server_default=text("CURRENT_TIMESTAMP"),
+        ),
+    )
+    metadata_: Optional[Dict[str, str]] = Field(default=None, sa_column=Column(JSON))

backend/api/quivr_api/modules/knowledge/repository/knowledge_interface.py

@@ -2,17 +2,15 @@
 from typing import List
 from uuid import UUID
 
-from quivr_api.modules.knowledge.dto.inputs import (
-    CreateKnowledgeProperties,
-    KnowledgeStatus,
-)
+from quivr_api.modules.knowledge.dto.inputs import (CreateKnowledgeProperties,
+                                                    KnowledgeStatus)
 from quivr_api.modules.knowledge.dto.outputs import DeleteKnowledgeResponse
-from quivr_core.models import QuivrKnowledge as Knowledge
+from quivr_api.modules.knowledge.entity.knowledge import KnowledgeDB
 
 
 class KnowledgeInterface(ABC):
     @abstractmethod
-    def insert_knowledge(self, knowledge: CreateKnowledgeProperties) -> Knowledge:
+    def insert_knowledge(self, knowledge: CreateKnowledgeProperties) -> KnowledgeDB:
         """
         Add a knowledge
         """
@@ -34,7 +32,7 @@ def remove_knowledge_by_id(
         pass
 
     @abstractmethod
-    def get_knowledge_by_id(self, knowledge_id: UUID) -> Knowledge:
+    def get_knowledge_by_id(self, knowledge_id: UUID) -> KnowledgeDB:
         """
         Get a knowledge by its id
         Args:
@@ -43,7 +41,7 @@ def get_knowledge_by_id(self, knowledge_id: UUID) -> Knowledge:
         pass
 
     @abstractmethod
-    def get_all_knowledge_in_brain(self, brain_id: UUID) -> List[Knowledge]:
+    def get_all_knowledge_in_brain(self, brain_id: UUID) -> List[KnowledgeDB]:
         """
         Get all the knowledge in a brain
         Args: