Skip to content

Commit

Permalink
feat(assistants): mock api (#3195)
Browse files Browse the repository at this point in the history
# Description

Please include a summary of the changes and the related issue. Please
also include relevant motivation and context.

## Checklist before requesting a review

Please delete options that are not relevant.

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented hard-to-understand areas
- [ ] I have ideally added tests that prove my fix is effective or that
my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged

## Screenshots (if appropriate):
  • Loading branch information
StanGirard authored Sep 18, 2024
1 parent 4390d31 commit 282fa0e
Show file tree
Hide file tree
Showing 109 changed files with 922 additions and 882 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -102,3 +102,4 @@ backend/core/examples/chatbot/.chainlit/translations/en-US.json
# Tox
.tox
Pipfile
*.pkl
1 change: 0 additions & 1 deletion backend/api/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ dependencies = [
"pydantic-settings>=2.4.0",
"python-dotenv>=1.0.1",
"unidecode>=1.3.8",
"fpdf>=1.7.2",
"colorlog>=6.8.2",
"posthog>=3.5.0",
"pyinstrument>=4.7.2",
Expand Down
2 changes: 1 addition & 1 deletion backend/api/quivr_api/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from colorlog import (
ColoredFormatter,
) # You need to install this package: pip install colorlog
)


def get_logger(logger_name, log_file="application.log"):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from jose import jwt
from jose.exceptions import JWTError

from quivr_api.modules.user.entity.user_identity import UserIdentity

SECRET_KEY = os.environ.get("JWT_SECRET_KEY")
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from uuid import UUID

from pydantic import BaseModel, ConfigDict

from quivr_api.logger import get_logger

logger = get_logger(__name__)
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from uuid import UUID

from fastapi import APIRouter, Depends, Query

from quivr_api.middlewares.auth.auth_bearer import AuthBearer, get_current_user
from quivr_api.modules.analytics.entity.analytics import Range
from quivr_api.modules.analytics.service.analytics_service import AnalyticsService
Expand Down
8 changes: 6 additions & 2 deletions backend/api/quivr_api/modules/analytics/entity/analytics.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,20 @@
from datetime import date
from enum import IntEnum
from typing import List

from pydantic import BaseModel
from datetime import date


class Range(IntEnum):
WEEK = 7
MONTH = 30
QUARTER = 90


class Usage(BaseModel):
date: date
usage_count: int


class BrainsUsages(BaseModel):
usages: List[Usage]
usages: List[Usage]
Original file line number Diff line number Diff line change
Expand Up @@ -11,5 +11,4 @@ def __init__(self):
self.repository = Analytics()

def get_brains_usages(self, user_id, graph_range, brain_id=None):

return self.repository.get_brains_usages(user_id, graph_range, brain_id)
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from uuid import uuid4

from fastapi import APIRouter, Depends

from quivr_api.logger import get_logger
from quivr_api.middlewares.auth import AuthBearer, get_current_user
from quivr_api.modules.api_key.dto.outputs import ApiKeyInfo
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from datetime import datetime

from fastapi import HTTPException

from quivr_api.logger import get_logger
from quivr_api.modules.api_key.repository.api_key_interface import ApiKeysInterface
from quivr_api.modules.api_key.repository.api_keys import ApiKeys
Expand Down
Original file line number Diff line number Diff line change
@@ -1 +1,6 @@
# noqa:
from .assistant_routes import assistant_router

__all__ = [
"assistant_router",
]
193 changes: 153 additions & 40 deletions backend/api/quivr_api/modules/assistant/controller/assistant_routes.py
Original file line number Diff line number Diff line change
@@ -1,63 +1,176 @@
from typing import List
import io
from typing import Annotated, List
from uuid import uuid4

from fastapi import APIRouter, Depends, HTTPException, UploadFile
from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile

from quivr_api.celery_config import celery
from quivr_api.logger import get_logger
from quivr_api.middlewares.auth import AuthBearer, get_current_user
from quivr_api.modules.assistant.dto.inputs import InputAssistant
from quivr_api.middlewares.auth.auth_bearer import AuthBearer, get_current_user
from quivr_api.modules.assistant.controller.assistants_definition import (
assistants,
validate_assistant_input,
)
from quivr_api.modules.assistant.dto.inputs import CreateTask, InputAssistant
from quivr_api.modules.assistant.dto.outputs import AssistantOutput
from quivr_api.modules.assistant.ito.difference import DifferenceAssistant
from quivr_api.modules.assistant.ito.summary import SummaryAssistant, summary_inputs
from quivr_api.modules.assistant.service.assistant import Assistant
from quivr_api.modules.assistant.entity.assistant_entity import (
AssistantSettings,
)
from quivr_api.modules.assistant.services.tasks_service import TasksService
from quivr_api.modules.dependencies import get_service
from quivr_api.modules.upload.service.upload_file import (
upload_file_storage,
)
from quivr_api.modules.user.entity.user_identity import UserIdentity

assistant_router = APIRouter()
logger = get_logger(__name__)

assistant_service = Assistant()

assistant_router = APIRouter()


TasksServiceDep = Annotated[TasksService, Depends(get_service(TasksService))]
UserIdentityDep = Annotated[UserIdentity, Depends(get_current_user)]


@assistant_router.get(
"/assistants", dependencies=[Depends(AuthBearer())], tags=["Assistant"]
)
async def list_assistants(
async def get_assistants(
request: Request,
current_user: UserIdentity = Depends(get_current_user),
) -> List[AssistantOutput]:
"""
Retrieve and list all the knowledge in a brain.
"""
logger.info("Getting assistants")

return assistants

summary = summary_inputs()
# difference = difference_inputs()
# crawler = crawler_inputs()
return [summary]

@assistant_router.get(
"/assistants/tasks", dependencies=[Depends(AuthBearer())], tags=["Assistant"]
)
async def get_tasks(
request: Request,
current_user: UserIdentityDep,
tasks_service: TasksServiceDep,
):
logger.info("Getting tasks")
return await tasks_service.get_tasks_by_user_id(current_user.id)


@assistant_router.post(
"/assistant/process",
dependencies=[Depends(AuthBearer())],
tags=["Assistant"],
"/assistants/task", dependencies=[Depends(AuthBearer())], tags=["Assistant"]
)
async def process_assistant(
async def create_task(
current_user: UserIdentityDep,
tasks_service: TasksServiceDep,
request: Request,
input: InputAssistant,
files: List[UploadFile] = None,
current_user: UserIdentity = Depends(get_current_user),
):
if input.name.lower() == "summary":
summary_assistant = SummaryAssistant(
input=input, files=files, current_user=current_user
)
try:
summary_assistant.check_input()
return await summary_assistant.process_assistant()
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
elif input.name.lower() == "difference":
difference_assistant = DifferenceAssistant(
input=input, files=files, current_user=current_user
)
assistant = next(
(assistant for assistant in assistants if assistant.id == input.id), None
)
if assistant is None:
raise HTTPException(status_code=404, detail="Assistant not found")

is_valid, validation_errors = validate_assistant_input(input, assistant)
if not is_valid:
for error in validation_errors:
print(error)
raise HTTPException(status_code=400, detail=error)
else:
print("Assistant input is valid.")
notification_uuid = uuid4()

# Process files dynamically
for upload_file in files:
file_name_path = f"{input.id}/{notification_uuid}/{upload_file.filename}"
buff_reader = io.BufferedReader(upload_file.file) # type: ignore
try:
difference_assistant.check_input()
return await difference_assistant.process_assistant()
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
return {"message": "Assistant not found"}
await upload_file_storage(buff_reader, file_name_path)
except Exception as e:
logger.exception(f"Exception in upload_route {e}")
raise HTTPException(
status_code=500, detail=f"Failed to upload file to storage. {e}"
)

task = CreateTask(
assistant_id=input.id,
pretty_id=str(notification_uuid),
settings=input.model_dump(mode="json"),
)

task_created = await tasks_service.create_task(task, current_user.id)

celery.send_task(
"process_assistant_task",
kwargs={
"assistant_id": input.id,
"notification_uuid": notification_uuid,
"task_id": task_created.id,
"user_id": str(current_user.id),
},
)
return task_created


@assistant_router.get(
"/assistants/task/{task_id}",
dependencies=[Depends(AuthBearer())],
tags=["Assistant"],
)
async def get_task(
request: Request,
task_id: str,
current_user: UserIdentityDep,
tasks_service: TasksServiceDep,
):
return await tasks_service.get_task_by_id(task_id, current_user.id) # type: ignore


@assistant_router.delete(
"/assistants/task/{task_id}",
dependencies=[Depends(AuthBearer())],
tags=["Assistant"],
)
async def delete_task(
request: Request,
task_id: int,
current_user: UserIdentityDep,
tasks_service: TasksServiceDep,
):
return await tasks_service.delete_task(task_id, current_user.id)


@assistant_router.get(
"/assistants/task/{task_id}/download",
dependencies=[Depends(AuthBearer())],
tags=["Assistant"],
)
async def get_download_link_task(
request: Request,
task_id: int,
current_user: UserIdentityDep,
tasks_service: TasksServiceDep,
):
return await tasks_service.get_download_link_task(task_id, current_user.id)


@assistant_router.get(
"/assistants/{assistant_id}/config",
dependencies=[Depends(AuthBearer())],
tags=["Assistant"],
response_model=AssistantSettings,
summary="Retrieve assistant configuration",
description="Get the settings and file requirements for the specified assistant.",
)
async def get_assistant_config(
assistant_id: int,
current_user: UserIdentityDep,
):
assistant = next(
(assistant for assistant in assistants if assistant.id == assistant_id), None
)
if assistant is None:
raise HTTPException(status_code=404, detail="Assistant not found")
return assistant.settings
Loading

0 comments on commit 282fa0e

Please sign in to comment.