256 lines
9.7 KiB
Python
256 lines
9.7 KiB
Python
import os
|
|
import tempfile
|
|
import uuid
|
|
from functools import partial
|
|
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter, BackgroundTasks, Body, Depends, HTTPException, Query, UploadFile, status
|
|
from fastapi.responses import JSONResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
from memgpt.data_sources.connectors import DirectoryConnector
|
|
from memgpt.data_types import Source
|
|
from memgpt.models.pydantic_models import DocumentModel, JobModel, JobStatus, PassageModel, SourceModel
|
|
from memgpt.server.rest_api.auth_token import get_current_user
|
|
from memgpt.server.rest_api.interface import QueuingInterface
|
|
from memgpt.server.server import SyncServer
|
|
|
|
router = APIRouter()
|
|
|
|
"""
|
|
Implement the following functions:
|
|
* List all available sources
|
|
* Create a new source
|
|
* Delete a source
|
|
* Upload a file to a server that is loaded into a specific source
|
|
* Paginated get all passages from a source
|
|
* Paginated get all documents from a source
|
|
* Attach a source to an agent
|
|
"""
|
|
|
|
|
|
class ListSourcesResponse(BaseModel):
|
|
sources: List[SourceModel] = Field(..., description="List of available sources.")
|
|
|
|
|
|
class CreateSourceRequest(BaseModel):
|
|
name: str = Field(..., description="The name of the source.")
|
|
description: Optional[str] = Field(None, description="The description of the source.")
|
|
|
|
|
|
class UploadFileToSourceRequest(BaseModel):
|
|
file: UploadFile = Field(..., description="The file to upload.")
|
|
|
|
|
|
class UploadFileToSourceResponse(BaseModel):
|
|
source: SourceModel = Field(..., description="The source the file was uploaded to.")
|
|
added_passages: int = Field(..., description="The number of passages added to the source.")
|
|
added_documents: int = Field(..., description="The number of documents added to the source.")
|
|
|
|
|
|
class GetSourcePassagesResponse(BaseModel):
|
|
passages: List[PassageModel] = Field(..., description="List of passages from the source.")
|
|
|
|
|
|
class GetSourceDocumentsResponse(BaseModel):
|
|
documents: List[DocumentModel] = Field(..., description="List of documents from the source.")
|
|
|
|
|
|
def load_file_to_source(server: SyncServer, user_id: uuid.UUID, source: Source, job_id: uuid.UUID, file: UploadFile):
|
|
# update job status
|
|
job = server.ms.get_job(job_id=job_id)
|
|
job.status = JobStatus.running
|
|
server.ms.update_job(job)
|
|
|
|
try:
|
|
# write the file to a temporary directory (deleted after the context manager exits)
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
file_path = os.path.join(tmpdirname, file.filename)
|
|
with open(file_path, "wb") as buffer:
|
|
buffer.write(file.file.read())
|
|
|
|
# read the file
|
|
connector = DirectoryConnector(input_files=[file_path])
|
|
|
|
# TODO: pre-compute total number of passages?
|
|
|
|
# load the data into the source via the connector
|
|
num_passages, num_documents = server.load_data(user_id=user_id, source_name=source.name, connector=connector)
|
|
except Exception as e:
|
|
# job failed with error
|
|
error = str(e)
|
|
print(error)
|
|
job.status = JobStatus.failed
|
|
job.metadata_["error"] = error
|
|
server.ms.update_job(job)
|
|
# TODO: delete any associated passages/documents?
|
|
return 0, 0
|
|
|
|
# update job status
|
|
job.status = JobStatus.completed
|
|
job.metadata_["num_passages"] = num_passages
|
|
job.metadata_["num_documents"] = num_documents
|
|
print("job completed", job.metadata_, job.id)
|
|
server.ms.update_job(job)
|
|
|
|
|
|
def setup_sources_index_router(server: SyncServer, interface: QueuingInterface, password: str):
|
|
get_current_user_with_server = partial(partial(get_current_user, server), password)
|
|
|
|
@router.get("/sources", tags=["sources"], response_model=ListSourcesResponse)
|
|
async def list_sources(
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
List all data sources created by a user.
|
|
"""
|
|
# Clear the interface
|
|
interface.clear()
|
|
|
|
try:
|
|
sources = server.list_all_sources(user_id=user_id)
|
|
return ListSourcesResponse(sources=sources)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"{e}")
|
|
|
|
@router.post("/sources", tags=["sources"], response_model=SourceModel)
|
|
async def create_source(
|
|
request: CreateSourceRequest = Body(...),
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Create a new data source.
|
|
"""
|
|
interface.clear()
|
|
try:
|
|
# TODO: don't use Source and just use SourceModel once pydantic migration is complete
|
|
source = server.create_source(name=request.name, user_id=user_id)
|
|
return SourceModel(
|
|
name=source.name,
|
|
description=None, # TODO: actually store descriptions
|
|
user_id=source.user_id,
|
|
id=source.id,
|
|
embedding_config=server.server_embedding_config,
|
|
created_at=source.created_at.timestamp(),
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"{e}")
|
|
|
|
@router.delete("/sources/{source_id}", tags=["sources"])
|
|
async def delete_source(
|
|
source_id: uuid.UUID,
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Delete a data source.
|
|
"""
|
|
interface.clear()
|
|
try:
|
|
server.delete_source(source_id=source_id, user_id=user_id)
|
|
return JSONResponse(status_code=status.HTTP_200_OK, content={"message": f"Source source_id={source_id} successfully deleted"})
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"{e}")
|
|
|
|
@router.post("/sources/{source_id}/attach", tags=["sources"], response_model=SourceModel)
|
|
async def attach_source_to_agent(
|
|
source_id: uuid.UUID,
|
|
agent_id: uuid.UUID = Query(..., description="The unique identifier of the agent to attach the source to."),
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Attach a data source to an existing agent.
|
|
"""
|
|
interface.clear()
|
|
assert isinstance(agent_id, uuid.UUID), f"Expected agent_id to be a UUID, got {agent_id}"
|
|
assert isinstance(user_id, uuid.UUID), f"Expected user_id to be a UUID, got {user_id}"
|
|
source = server.ms.get_source(source_id=source_id, user_id=user_id)
|
|
source = server.attach_source_to_agent(source_name=source.name, agent_id=agent_id, user_id=user_id)
|
|
return SourceModel(
|
|
name=source.name,
|
|
description=None, # TODO: actually store descriptions
|
|
user_id=source.user_id,
|
|
id=source.id,
|
|
embedding_config=server.server_embedding_config,
|
|
created_at=source.created_at,
|
|
)
|
|
|
|
@router.post("/sources/{source_id}/detach", tags=["sources"], response_model=SourceModel)
|
|
async def detach_source_from_agent(
|
|
source_id: uuid.UUID,
|
|
agent_id: uuid.UUID = Query(..., description="The unique identifier of the agent to detach the source from."),
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Detach a data source from an existing agent.
|
|
"""
|
|
server.detach_source_from_agent(source_id=source_id, agent_id=agent_id, user_id=user_id)
|
|
|
|
@router.get("/sources/status/{job_id}", tags=["sources"], response_model=JobModel)
|
|
async def get_job_status(
|
|
job_id: uuid.UUID,
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Get the status of a job.
|
|
"""
|
|
job = server.ms.get_job(job_id=job_id)
|
|
if job is None:
|
|
raise HTTPException(status_code=404, detail=f"Job with id={job_id} not found.")
|
|
return job
|
|
|
|
@router.post("/sources/{source_id}/upload", tags=["sources"], response_model=JobModel)
|
|
async def upload_file_to_source(
|
|
# file: UploadFile = UploadFile(..., description="The file to upload."),
|
|
file: UploadFile,
|
|
source_id: uuid.UUID,
|
|
background_tasks: BackgroundTasks,
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Upload a file to a data source.
|
|
"""
|
|
interface.clear()
|
|
source = server.ms.get_source(source_id=source_id, user_id=user_id)
|
|
|
|
# create job
|
|
job = JobModel(user_id=user_id, metadata={"type": "embedding", "filename": file.filename, "source_id": source_id})
|
|
job_id = job.id
|
|
server.ms.create_job(job)
|
|
|
|
# create background task
|
|
background_tasks.add_task(load_file_to_source, server, user_id, source, job_id, file)
|
|
|
|
# return job information
|
|
job = server.ms.get_job(job_id=job_id)
|
|
return job
|
|
|
|
@router.get("/sources/{source_id}/passages ", tags=["sources"], response_model=GetSourcePassagesResponse)
|
|
async def list_passages(
|
|
source_id: uuid.UUID,
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
List all passages associated with a data source.
|
|
"""
|
|
passages = server.list_data_source_passages(user_id=user_id, source_id=source_id)
|
|
return GetSourcePassagesResponse(passages=passages)
|
|
|
|
@router.get("/sources/{source_id}/documents", tags=["sources"], response_model=GetSourceDocumentsResponse)
|
|
async def list_documents(
|
|
source_id: uuid.UUID,
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
List all documents associated with a data source.
|
|
"""
|
|
documents = server.list_data_source_documents(user_id=user_id, source_id=source_id)
|
|
return GetSourceDocumentsResponse(documents=documents)
|
|
|
|
return router
|