add features

ocr/__init__.py

@@ -12,8 +12,11 @@ from ocr.core.wrappers import OcrResponseWrapper, ErrorOcrResponse
 def create_app() -> FastAPI:
     app = FastAPI()
 
-    from ocr.api.message import report_router
-    app.include_router(report_router, tags=['message'])
+    from ocr.api.report import report_router
+    app.include_router(report_router, tags=['report'])
+
+    from ocr.api.message import message_router
+    app.include_router(message_router, tags=['message'])
 
     app.add_middleware(
         CORSMiddleware,
@@ -41,6 +44,6 @@ def create_app() -> FastAPI:
 
     @app.get("/")
     async def read_root():
-        return {"message": "Hello world!"}
+        return {"report": "Hello world!"}
 
     return app

ocr/api/message/__init__.py

@@ -1,7 +1,7 @@
 from fastapi.routing import APIRouter
 
-report_router = APIRouter(
-    prefix="/api/report", tags=["message"]
+message_router = APIRouter(
+    prefix="/api/message", tags=["message"]
 )
 
 from . import views

ocr/api/message/ai/openai_request.py

@@ -1,17 +0,0 @@
-from ocr.api.message.ai.prompts import OCRPrompts
-from ocr.core.wrappers import openai_wrapper
-
-
-@openai_wrapper(model='gpt-4o-mini')
-async def generate_report(request_content: list[dict]):
-    messages = [
-        {
-            "role": "system",
-            "content": OCRPrompts.generate_general_answer
-        },
-        {
-            "role": "user",
-            "content": request_content
-        }
-    ]
-    return messages

ocr/api/message/ai/prompts.py

@@ -1,44 +0,0 @@
-class OCRPrompts:
-    generate_general_answer = """## Task
-
-You must analyze the text extracted from medical document and generate a comprehensive report in **Markdown2** format. Ensure that every detail provided in the document is included, and do not omit or modify any information. Your output must strictly follow the required format.
-
-## Report Structure
-
-The report should be structured as follows, with each section containing only relevant information from the document:
-
-```markdown
-## Patient Information
-
-- Name: [Patient Name]
-- Age: [Patient Age]
-- Date of Scan: [Date]
-- Indication: [Reason for the CT scan]
-
-## Findings
-
-**Primary findings**:
-[Describe significant abnormalities or findings relevant to the indication]
-
-** Secondary findings**:
-[List incidental findings, e.g., "Mild hepatic steatosis noted."]
-**No abnormalities**:
-[Mention organs or systems without abnormalities, e.g., "No evidence of lymphadenopathy or pleural effusion."]
-
-## Impression
-
-[Summarize the findings concisely, e.g., "Findings suggest a primary lung tumor. Biopsy recommended for further evaluation."]
-
-## Recommendations
-
-[Include next steps or further tests, e.g., "PET scan and consultation with oncology recommended."]
-```
-
-[INST]
-
-## Instructions
-
-- **Do not invent or infer any information.** Only use data provided in the user request.
-- Ensure that the format is followed strictly, and the output is complete without any deviations.
-
-[/INST]"""

ocr/api/message/db_requests.py

@@ -1,26 +1,39 @@
 import asyncio
 
-from ocr.api.message.model import MessageModel
-from ocr.core.config import settings
+from fastapi import HTTPException
 
+from ocr.api.message.dto import Author
+from ocr.api.message.models import MessageModel
+from ocr.api.message.schemas import CreateMessageRequest
+from ocr.core.config import settings
 
-async def get_all_chat_messages_obj(page_size: int, page_index: int) -> tuple[list[MessageModel], int]:
-    skip = page_size * page_index
-    objects, total_count = await asyncio.gather(
-        settings.DB_CLIENT.messages
-        .find()
-        .skip(skip)
-        .limit(page_size)
-        .to_list(length=page_size),
-        settings.DB_CLIENT.messages.count_documents({})
-    )
-    return objects, total_count
 
+async def create_message_obj(
+        report_id: str, message_data: CreateMessageRequest
+) -> MessageModel:
+    report = await settings.DB_CLIENT.reports.find_one({"id": report_id})
+    if not report:
+        raise HTTPException(status_code=404, detail="Report not found.")
 
-async def save_report_obj(report: str, filename: str) -> MessageModel:
-    message = MessageModel(
-        text=report,
-        filename=filename,
-    )
+    message = MessageModel(**message_data.model_dump(), reportId=report_id, author=Author.User)
     await settings.DB_CLIENT.messages.insert_one(message.to_mongo())
     return message
+
+
+async def get_all_chat_messages_obj(report_id: str) -> list[MessageModel]:
+    messages, report = await asyncio.gather(
+        settings.DB_CLIENT.messages.find({"reportId": report_id}).to_list(length=None),
+        settings.DB_CLIENT.reports.find_one({"id": report_id})
+    )
+    messages = [MessageModel.from_mongo(message) for message in messages]
+    if not report:
+        raise HTTPException(status_code=404, detail="Report not found")
+    return messages
+
+
+async def save_assistant_user_message(user_message: str, assistant_message: str, report_id: str) -> MessageModel:
+    user_message = MessageModel(reportId=report_id, author=Author.User, text=user_message)
+    assistant_message = MessageModel(reportId=report_id, author=Author.Assistant, text=assistant_message)
+    await settings.DB_CLIENT.messages.insert_one(user_message.to_mongo())
+    await settings.DB_CLIENT.messages.insert_one(assistant_message.to_mongo())
+    return user_message

ocr/api/message/dto.py

@@ -1,7 +1,8 @@
+from enum import Enum
+
 from pydantic import BaseModel
 
 
-class Paging(BaseModel):
-    pageSize: int
-    pageIndex: int
-    totalCount: int
+class Author(Enum):
+    User = "user"
+    Assistant = "assistant"

ocr/api/message/models.py

@@ -0,0 +1,14 @@
+from datetime import datetime
+
+from pydantic import Field
+
+from ocr.api.message.dto import Author
+from ocr.core.database import MongoBaseModel
+
+
+class MessageModel(MongoBaseModel):
+    reportId: str
+    author: Author
+    text: str
+    datetimeInserted: datetime = Field(default_factory=datetime.now)
+    datetimeUpdated: datetime = Field(default_factory=datetime.now)

ocr/api/message/schemas.py

@@ -1,14 +1,13 @@
 from pydantic import BaseModel
 
-from ocr.api.message.dto import Paging
-from ocr.api.message.model import MessageModel
+from ocr.api.message.models import MessageModel
+from ocr.api.report.dto import Paging
 from ocr.core.wrappers import OcrResponseWrapper
 
 
 class CreateMessageRequest(BaseModel):
     text: str
 
-
 class MessageWrapper(OcrResponseWrapper[MessageModel]):
     pass
 

ocr/api/message/views.py

@@ -1,24 +1,17 @@
-from typing import Optional
-
-from fastapi import Query, UploadFile, File
-
-from ocr.api.message import report_router
-from ocr.api.message.ai.openai_request import generate_report
-from ocr.api.message.db_requests import get_all_chat_messages_obj, save_report_obj
-from ocr.api.message.dto import Paging
-from ocr.api.message.model import MessageModel
-from ocr.api.message.schemas import (AllMessageWrapper,
-                                     AllMessageResponse)
-from ocr.api.message.utils import divide_images, prepare_request_content, clean_response
+from ocr.api.message import message_router
+from ocr.api.message.db_requests import get_all_chat_messages_obj, save_assistant_user_message
+from ocr.api.message.models import MessageModel
+from ocr.api.message.schemas import AllMessageWrapper, AllMessageResponse, CreateMessageRequest
+from ocr.api.report.dto import Paging
+from ocr.api.utils import transform_messages_to_openai
 from ocr.core.wrappers import OcrResponseWrapper
 
 
-@report_router.get('/all')
+@message_router.get('/{reportId}/all')
 async def get_all_chat_messages(
-        pageSize: Optional[int] = Query(10, description="Number of countries to return per page"),
-        pageIndex: Optional[int] = Query(0, description="Page index to retrieve"),
+        reportId: str
 ) -> AllMessageWrapper:
-    messages, _ = await get_all_chat_messages_obj(pageSize, pageIndex)
+    messages = await get_all_chat_messages_obj(reportId)
     response = AllMessageResponse(
         paging=Paging(pageSize=len(messages), pageIndex=0, totalCount=len(messages)),
         data=messages
@@ -26,16 +19,14 @@ async def get_all_chat_messages(
     return AllMessageWrapper(data=response)
 
 
-@report_router.post('/generate')
+@message_router.post('/{reportId}')
 async def create_message(
-        file: UploadFile = File(...),
+        reportId: str,
+        message_data: CreateMessageRequest,
 ) -> OcrResponseWrapper[MessageModel]:
-    try:
-        contents = await file.read()
-        images = divide_images(contents)
-        content = prepare_request_content(images)
-        report = await generate_report(content)
-        response = await save_report_obj(clean_response(report), file.filename)
-        return OcrResponseWrapper(data=response)
-    finally:
-        await file.close()
+    messages = await get_all_chat_messages_obj(reportId)
+    message_history = transform_messages_to_openai(messages)
+    # response = await generate_response()
+    response = 'Hello world'
+    response = await save_assistant_user_message(message_data.text, response, reportId)
+    return OcrResponseWrapper(data=response)

ocr/api/report/__init__.py

@@ -0,0 +1,7 @@
+from fastapi.routing import APIRouter
+
+report_router = APIRouter(
+    prefix="/api/report", tags=["report"]
+)
+
+from . import views

ocr/api/report/db_requests.py

@@ -0,0 +1,27 @@
+from fastapi import HTTPException
+
+from ocr.api.report.dto import ReportModelShort
+from ocr.api.report.model import ReportModel
+from ocr.core.config import settings
+
+
+async def get_all_reports_obj() -> list[ReportModelShort]:
+    reports = await settings.DB_CLIENT.reports.find({}).to_list(length=None)
+    return [ReportModelShort(**report) for report in reports]
+
+
+async def delete_all_reports() -> None:
+    await settings.DB_CLIENT.reports.delete_many({})
+
+
+async def get_report_obj_by_id(report_id: str) -> ReportModel:
+    report = await settings.DB_CLIENT.reports.find_one({"id": report_id})
+    if not report:
+        raise HTTPException(status_code=404, detail="Report not found")
+    return ReportModel.from_mongo(report)
+
+
+async def save_report_obj(report: str, changes: str) -> ReportModel:
+    report = ReportModel(report=report, changes=changes, filename='maksim.docx')
+    await settings.DB_CLIENT.reports.insert_one(report.to_mongo())
+    return report

ocr/api/report/dto.py

@@ -0,0 +1,16 @@
+from typing import ClassVar
+
+from pydantic import BaseModel
+
+from ocr.api.report.model import ReportModel
+
+
+class Paging(BaseModel):
+    pageSize: int
+    pageIndex: int
+    totalCount: int
+
+
+class ReportModelShort(ReportModel):
+    report: ClassVar[str]
+    changes: ClassVar[str]

ocr/api/{message → report}/model.py

@@ -5,8 +5,9 @@ from pydantic import Field
 from ocr.core.database import MongoBaseModel
 
 
-class MessageModel(MongoBaseModel):
-    text: str
+class ReportModel(MongoBaseModel):
+    report: str
+    changes: str
     filename: str
     datetimeInserted: datetime = Field(default_factory=datetime.now)
-    datetimeUpdated: datetime = Field(default_factory=datetime.now)
+    datetimeUpdated: datetime = Field(default_factory=datetime.now)

ocr/api/report/schemas.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+from ocr.api.report.dto import Paging, ReportModelShort
+
+
+class AllReportResponse(BaseModel):
+    paging: Paging
+    data: list[ReportModelShort]

ocr/api/report/views.py

@@ -0,0 +1,41 @@
+from fastapi import UploadFile, File
+
+from ocr.api.report import report_router
+from ocr.api.report.db_requests import get_all_reports_obj, delete_all_reports, get_report_obj_by_id, save_report_obj
+from ocr.api.report.dto import Paging
+from ocr.api.report.model import ReportModel
+from ocr.api.report.schemas import AllReportResponse
+from ocr.core.wrappers import OcrResponseWrapper
+
+
+@report_router.get('/all')
+async def get_all_reports() -> OcrResponseWrapper[AllReportResponse]:
+    reports = await get_all_reports_obj()
+    response = AllReportResponse(
+        paging=Paging(pageSize=len(reports), pageIndex=0, totalCount=len(reports)),
+        data=reports
+    )
+    return OcrResponseWrapper(data=response)
+
+
+@report_router.delete('/all')
+async def delete_all_report() -> OcrResponseWrapper:
+    await delete_all_reports()
+    return OcrResponseWrapper()
+
+
+@report_router.get('/{reportId}')
+async def get_report(reportId: str) -> OcrResponseWrapper[ReportModel]:
+    report = await get_report_obj_by_id(reportId)
+    return OcrResponseWrapper(data=report)
+
+
+@report_router.post('')
+async def create_report(
+        file: UploadFile = File(...),
+) -> OcrResponseWrapper[ReportModel]:
+    # messages = await create_new_reports(reportId)
+    # response = await generate_response(message_data.text, message_history)
+    report, changes = 'Hello', 'World'
+    report = await save_report_obj(report, changes)
+    return OcrResponseWrapper(data=report)

ocr/api/{message/utils.py → utils.py}

@@ -6,6 +6,20 @@ import pytesseract
 from PIL import Image
 from pdf2image import convert_from_bytes
 
+from ocr.api.message.models import MessageModel
+
+
+def transform_messages_to_openai(messages: list[MessageModel]) -> list[dict]:
+    openai_messages = []
+    for message in messages:
+        content = message.text
+        openai_messages.append({
+            "role": message.author.value,
+            "content": content
+        })
+
+    return openai_messages
+
 
 def divide_images(contents: bytes) -> list[bytes]:
     images = convert_from_bytes(contents, dpi=250)