mirror of
https://github.com/MODSetter/SurfSense.git
synced 2025-09-10 14:28:57 +00:00
Fixed all ruff lint and formatting errors
This commit is contained in:
Utkarsh-Patel-13
parent
0a03c42cc5
commit
d359a59f6d
85 changed files with 5520 additions and 3870 deletions
|
|
@ -1,38 +1,40 @@
|
|||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from langchain.schema import AIMessage, HumanMessage
|
||||
from sqlalchemy.exc import IntegrityError, OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.db import Chat, SearchSpace, User, get_async_session
|
||||
from app.schemas import AISDKChatRequest, ChatCreate, ChatRead, ChatUpdate
|
||||
from app.tasks.stream_connector_search_results import stream_connector_search_results
|
||||
from app.users import current_active_user
|
||||
from app.utils.check_ownership import check_ownership
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy.exc import IntegrityError, OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from langchain.schema import HumanMessage, AIMessage
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/chat")
|
||||
async def handle_chat_data(
|
||||
request: AISDKChatRequest,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user)
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
messages = request.messages
|
||||
if messages[-1]['role'] != "user":
|
||||
if messages[-1]["role"] != "user":
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Last message must be a user message")
|
||||
status_code=400, detail="Last message must be a user message"
|
||||
)
|
||||
|
||||
user_query = messages[-1]['content']
|
||||
search_space_id = request.data.get('search_space_id')
|
||||
research_mode: str = request.data.get('research_mode')
|
||||
selected_connectors: List[str] = request.data.get('selected_connectors')
|
||||
document_ids_to_add_in_context: List[int] = request.data.get('document_ids_to_add_in_context')
|
||||
|
||||
search_mode_str = request.data.get('search_mode', "CHUNKS")
|
||||
user_query = messages[-1]["content"]
|
||||
search_space_id = request.data.get("search_space_id")
|
||||
research_mode: str = request.data.get("research_mode")
|
||||
selected_connectors: list[str] = request.data.get("selected_connectors")
|
||||
document_ids_to_add_in_context: list[int] = request.data.get(
|
||||
"document_ids_to_add_in_context"
|
||||
)
|
||||
|
||||
search_mode_str = request.data.get("search_mode", "CHUNKS")
|
||||
|
||||
# Convert search_space_id to integer if it's a string
|
||||
if search_space_id and isinstance(search_space_id, str):
|
||||
|
|
@ -40,21 +42,23 @@ async def handle_chat_data(
|
|||
search_space_id = int(search_space_id)
|
||||
except ValueError:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Invalid search_space_id format")
|
||||
status_code=400, detail="Invalid search_space_id format"
|
||||
) from None
|
||||
|
||||
# Check if the search space belongs to the current user
|
||||
try:
|
||||
await check_ownership(session, SearchSpace, search_space_id, user)
|
||||
except HTTPException:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="You don't have access to this search space")
|
||||
|
||||
status_code=403, detail="You don't have access to this search space"
|
||||
) from None
|
||||
|
||||
langchain_chat_history = []
|
||||
for message in messages[:-1]:
|
||||
if message['role'] == "user":
|
||||
langchain_chat_history.append(HumanMessage(content=message['content']))
|
||||
elif message['role'] == "assistant":
|
||||
langchain_chat_history.append(AIMessage(content=message['content']))
|
||||
if message["role"] == "user":
|
||||
langchain_chat_history.append(HumanMessage(content=message["content"]))
|
||||
elif message["role"] == "assistant":
|
||||
langchain_chat_history.append(AIMessage(content=message["content"]))
|
||||
|
||||
response = StreamingResponse(
|
||||
stream_connector_search_results(
|
||||
|
|
@ -69,7 +73,7 @@ async def handle_chat_data(
|
|||
document_ids_to_add_in_context,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
response.headers["x-vercel-ai-data-stream"] = "v1"
|
||||
return response
|
||||
|
||||
|
|
@ -78,7 +82,7 @@ async def handle_chat_data(
|
|||
async def create_chat(
|
||||
chat: ChatCreate,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user)
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
try:
|
||||
await check_ownership(session, SearchSpace, chat.search_space_id, user)
|
||||
|
|
@ -89,52 +93,57 @@ async def create_chat(
|
|||
return db_chat
|
||||
except HTTPException:
|
||||
raise
|
||||
except IntegrityError as e:
|
||||
except IntegrityError:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Database constraint violation. Please check your input data.")
|
||||
except OperationalError as e:
|
||||
status_code=400,
|
||||
detail="Database constraint violation. Please check your input data.",
|
||||
) from None
|
||||
except OperationalError:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=503, detail="Database operation failed. Please try again later.")
|
||||
except Exception as e:
|
||||
status_code=503, detail="Database operation failed. Please try again later."
|
||||
) from None
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=500, detail="An unexpected error occurred while creating the chat.")
|
||||
status_code=500,
|
||||
detail="An unexpected error occurred while creating the chat.",
|
||||
) from None
|
||||
|
||||
|
||||
@router.get("/chats/", response_model=List[ChatRead])
|
||||
@router.get("/chats/", response_model=list[ChatRead])
|
||||
async def read_chats(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
search_space_id: int = None,
|
||||
search_space_id: int | None = None,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user)
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
try:
|
||||
query = select(Chat).join(SearchSpace).filter(SearchSpace.user_id == user.id)
|
||||
|
||||
|
||||
# Filter by search_space_id if provided
|
||||
if search_space_id is not None:
|
||||
query = query.filter(Chat.search_space_id == search_space_id)
|
||||
|
||||
result = await session.execute(
|
||||
query.offset(skip).limit(limit)
|
||||
)
|
||||
|
||||
result = await session.execute(query.offset(skip).limit(limit))
|
||||
return result.scalars().all()
|
||||
except OperationalError:
|
||||
raise HTTPException(
|
||||
status_code=503, detail="Database operation failed. Please try again later.")
|
||||
status_code=503, detail="Database operation failed. Please try again later."
|
||||
) from None
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code=500, detail="An unexpected error occurred while fetching chats.")
|
||||
status_code=500, detail="An unexpected error occurred while fetching chats."
|
||||
) from None
|
||||
|
||||
|
||||
@router.get("/chats/{chat_id}", response_model=ChatRead)
|
||||
async def read_chat(
|
||||
chat_id: int,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user)
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
try:
|
||||
result = await session.execute(
|
||||
|
|
@ -145,14 +154,19 @@ async def read_chat(
|
|||
chat = result.scalars().first()
|
||||
if not chat:
|
||||
raise HTTPException(
|
||||
status_code=404, detail="Chat not found or you don't have permission to access it")
|
||||
status_code=404,
|
||||
detail="Chat not found or you don't have permission to access it",
|
||||
)
|
||||
return chat
|
||||
except OperationalError:
|
||||
raise HTTPException(
|
||||
status_code=503, detail="Database operation failed. Please try again later.")
|
||||
status_code=503, detail="Database operation failed. Please try again later."
|
||||
) from None
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code=500, detail="An unexpected error occurred while fetching the chat.")
|
||||
status_code=500,
|
||||
detail="An unexpected error occurred while fetching the chat.",
|
||||
) from None
|
||||
|
||||
|
||||
@router.put("/chats/{chat_id}", response_model=ChatRead)
|
||||
|
|
@ -160,7 +174,7 @@ async def update_chat(
|
|||
chat_id: int,
|
||||
chat_update: ChatUpdate,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user)
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
try:
|
||||
db_chat = await read_chat(chat_id, session, user)
|
||||
|
|
@ -175,22 +189,27 @@ async def update_chat(
|
|||
except IntegrityError:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Database constraint violation. Please check your input data.")
|
||||
status_code=400,
|
||||
detail="Database constraint violation. Please check your input data.",
|
||||
) from None
|
||||
except OperationalError:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=503, detail="Database operation failed. Please try again later.")
|
||||
status_code=503, detail="Database operation failed. Please try again later."
|
||||
) from None
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=500, detail="An unexpected error occurred while updating the chat.")
|
||||
status_code=500,
|
||||
detail="An unexpected error occurred while updating the chat.",
|
||||
) from None
|
||||
|
||||
|
||||
@router.delete("/chats/{chat_id}", response_model=dict)
|
||||
async def delete_chat(
|
||||
chat_id: int,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user)
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
try:
|
||||
db_chat = await read_chat(chat_id, session, user)
|
||||
|
|
@ -202,81 +221,16 @@ async def delete_chat(
|
|||
except IntegrityError:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Cannot delete chat due to existing dependencies.")
|
||||
status_code=400, detail="Cannot delete chat due to existing dependencies."
|
||||
) from None
|
||||
except OperationalError:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=503, detail="Database operation failed. Please try again later.")
|
||||
status_code=503, detail="Database operation failed. Please try again later."
|
||||
) from None
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise HTTPException(
|
||||
status_code=500, detail="An unexpected error occurred while deleting the chat.")
|
||||
|
||||
|
||||
# test_data = [
|
||||
# {
|
||||
# "type": "TERMINAL_INFO",
|
||||
# "content": [
|
||||
# {
|
||||
# "id": 1,
|
||||
# "text": "Starting to search for crawled URLs...",
|
||||
# "type": "info"
|
||||
# },
|
||||
# {
|
||||
# "id": 2,
|
||||
# "text": "Found 2 relevant crawled URLs",
|
||||
# "type": "success"
|
||||
# }
|
||||
# ]
|
||||
# },
|
||||
# {
|
||||
# "type": "SOURCES",
|
||||
# "content": [
|
||||
# {
|
||||
# "id": 1,
|
||||
# "name": "Crawled URLs",
|
||||
# "type": "CRAWLED_URL",
|
||||
# "sources": [
|
||||
# {
|
||||
# "id": 1,
|
||||
# "title": "Webpage Title",
|
||||
# "description": "Webpage Dec",
|
||||
# "url": "https://jsoneditoronline.org/"
|
||||
# },
|
||||
# {
|
||||
# "id": 2,
|
||||
# "title": "Webpage Title",
|
||||
# "description": "Webpage Dec",
|
||||
# "url": "https://www.google.com/"
|
||||
# }
|
||||
# ]
|
||||
# },
|
||||
# {
|
||||
# "id": 2,
|
||||
# "name": "Files",
|
||||
# "type": "FILE",
|
||||
# "sources": [
|
||||
# {
|
||||
# "id": 3,
|
||||
# "title": "Webpage Title",
|
||||
# "description": "Webpage Dec",
|
||||
# "url": "https://jsoneditoronline.org/"
|
||||
# },
|
||||
# {
|
||||
# "id": 4,
|
||||
# "title": "Webpage Title",
|
||||
# "description": "Webpage Dec",
|
||||
# "url": "https://www.google.com/"
|
||||
# }
|
||||
# ]
|
||||
# }
|
||||
# ]
|
||||
# },
|
||||
# {
|
||||
# "type": "ANSWER",
|
||||
# "content": [
|
||||
# "## SurfSense Introduction",
|
||||
# "Surfsense is A Personal NotebookLM and Perplexity-like AI Assistant for Everyone. Research and Never forget Anything. [1] [3]"
|
||||
# ]
|
||||
# }
|
||||
# ]
|
||||
status_code=500,
|
||||
detail="An unexpected error occurred while deleting the chat.",
|
||||
) from None
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue