60 lines
No EOL
2.3 KiB
Python
60 lines
No EOL
2.3 KiB
Python
from fastapi import Form, UploadFile, Depends
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi import APIRouter
|
|
from typing import List
|
|
import aiohttp
|
|
import subprocess
|
|
import logging
|
|
import os
|
|
from cognee.modules.users.models import User
|
|
from cognee.modules.users.methods import get_authenticated_user
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def get_add_router() -> APIRouter:
|
|
router = APIRouter()
|
|
|
|
@router.post("/", response_model=None)
|
|
async def add(
|
|
data: List[UploadFile],
|
|
datasetId: str = Form(...),
|
|
user: User = Depends(get_authenticated_user),
|
|
):
|
|
""" This endpoint is responsible for adding data to the graph."""
|
|
from cognee.api.v1.add import add as cognee_add
|
|
try:
|
|
if isinstance(data, str) and data.startswith("http"):
|
|
if "github" in data:
|
|
# Perform git clone if the URL is from GitHub
|
|
repo_name = data.split("/")[-1].replace(".git", "")
|
|
subprocess.run(["git", "clone", data, f".data/{repo_name}"], check=True)
|
|
await cognee_add(
|
|
"data://.data/",
|
|
f"{repo_name}",
|
|
)
|
|
else:
|
|
# Fetch and store the data from other types of URL using curl
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(data) as resp:
|
|
if resp.status == 200:
|
|
file_data = await resp.read()
|
|
filename = os.path.basename(data)
|
|
with open(f".data/{filename}", "wb") as f:
|
|
f.write(file_data)
|
|
await cognee_add(
|
|
"data://.data/",
|
|
f"{data.split('/')[-1]}",
|
|
)
|
|
else:
|
|
await cognee_add(
|
|
data,
|
|
datasetId,
|
|
user=user,
|
|
)
|
|
except Exception as error:
|
|
return JSONResponse(
|
|
status_code=409,
|
|
content={"error": str(error)}
|
|
)
|
|
|
|
return router |