github-starred/flowsint
2
0
Fork 0
mirror of https://github.com/reconurge/flowsint.git synced 2026-05-05 11:18:55 -05:00
Code Issues 264 Packages Projects Releases 11 Wiki Activity

feat: python formatter finally working

Browse Source
This commit is contained in:
dextmorgn
2025-08-13 20:08:06 +02:00
parent 4d5f96bb8d
commit 7620a6d145
148 changed files with 6777 additions and 3269 deletions
Download Patch File Download Diff File Expand all files Collapse all files

161
flowsint-api/app/api/routes/chat.py
View File

@@ -14,6 +14,7 @@ from flowsint_core.core.postgre_db import get_db
from app.models.models import Chat, ChatMessage, Profile
from app.api.deps import get_current_user
from app.api.schemas.chat import ChatCreate, ChatRead
router = APIRouter()
@@ -25,22 +26,23 @@ def clean_context(context: List[Dict]) -> List[Dict]:
if isinstance(item, dict):
# Create a copy and remove unwanted keys
cleaned_item = item["data"].copy()
# Remove top-level keys
cleaned_item.pop('id', None)
cleaned_item.pop('sketch_id', None)
cleaned_item.pop("id", None)
cleaned_item.pop("sketch_id", None)
# Remove from data if it exists
if 'data' in cleaned_item and isinstance(cleaned_item['data'], dict):
cleaned_item['data'].pop('sketch_id', None)
if "data" in cleaned_item and isinstance(cleaned_item["data"], dict):
cleaned_item["data"].pop("sketch_id", None)
# Remove measured/dimensions
cleaned_item.pop('measured', None)
cleaned_item.pop("measured", None)
cleaned.append(cleaned_item)
print(cleaned)
return cleaned
class ChatRequest(BaseModel):
prompt: str
context: Optional[List[Dict]] = None
@@ -49,17 +51,14 @@ class ChatRequest(BaseModel):
# Get all chats
@router.get("/", response_model=List[ChatRead])
def get_chats(
db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user)
db: Session = Depends(get_db), current_user: Profile = Depends(get_current_user)
):
chats = db.query(Chat).filter(
Chat.owner_id == current_user.id
).all()
chats = db.query(Chat).filter(Chat.owner_id == current_user.id).all()
# Sort messages for each chat by created_at in ascending order
for chat in chats:
chat.messages.sort(key=lambda x: x.created_at)
return chats
@@ -69,7 +68,7 @@ def get_chats(
# for chat in chats:
# db.delete(chat)
# db.commit()
# return {"result": "done"}
# return {"result": "done"}
# Get analyses by investigation ID
@@ -77,26 +76,39 @@ def get_chats(
def get_chats_by_investigation(
investigation_id: UUID,
db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user)
current_user: Profile = Depends(get_current_user),
):
chats = db.query(Chat).filter(
Chat.investigation_id == investigation_id,
Chat.owner_id == current_user.id
).all()
chats = (
db.query(Chat)
.filter(
Chat.investigation_id == investigation_id, Chat.owner_id == current_user.id
)
.all()
)
# Sort messages for each chat by created_at in ascending order
for chat in chats:
chat.messages.sort(key=lambda x: x.created_at)
return chats
@router.post("/stream/{chat_id}")
async def stream_chat(chat_id: UUID, payload: ChatRequest, db: Session = Depends(get_db), current_user: Profile = Depends(get_current_user)):
async def stream_chat(
chat_id: UUID,
payload: ChatRequest,
db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user),
):
# Check if Chat exists
chat = db.query(Chat).filter(Chat.id == chat_id, Chat.owner_id == current_user.id).first()
chat = (
db.query(Chat)
.filter(Chat.id == chat_id, Chat.owner_id == current_user.id)
.first()
)
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
# Update chat's last_updated_at
chat.last_updated_at = datetime.utcnow()
db.commit()
@@ -112,19 +124,25 @@ async def stream_chat(chat_id: UUID, payload: ChatRequest, db: Session = Depends
db.add(user_message)
db.commit()
db.refresh(user_message)
try:
api_key = os.environ.get("MISTRAL_API_KEY")
if not api_key:
raise HTTPException(status_code=500, detail="Mistral API key not configured")
raise HTTPException(
status_code=500, detail="Mistral API key not configured"
)
client = Mistral(api_key=api_key)
model = "mistral-small-latest"
accumulated_content = []
context_message = None
# Convert database messages to Mistral format
messages = [SystemMessage(content="You are a CTI/OSINT investigator and you are trying to investigate on a variety of real life cases. Use your knowledge and analytics capabilities to analyse the context and answer the question the best you can. If you need to reference some items (an IP, a domain or something particular) please use the code brackets, like : `12.23.34.54` to reference it.")]
messages = [
SystemMessage(
content="You are a CTI/OSINT investigator and you are trying to investigate on a variety of real life cases. Use your knowledge and analytics capabilities to analyse the context and answer the question the best you can. If you need to reference some items (an IP, a domain or something particular) please use the code brackets, like : `12.23.34.54` to reference it."
)
]
# Add context as a single system message if provided
if payload.context:
try:
@@ -139,16 +157,22 @@ async def stream_chat(chat_id: UUID, payload: ChatRequest, db: Session = Depends
except Exception as e:
# If context processing fails, skip it
print(f"Context processing error: {e}")
# Sort messages by created_at in ascending order and get recent messages
sorted_messages = sorted(chat.messages, key=lambda x: x.created_at)
recent_messages = sorted_messages[-5:] if len(sorted_messages) > 5 else sorted_messages
recent_messages = (
sorted_messages[-5:] if len(sorted_messages) > 5 else sorted_messages
)
for message in recent_messages:
if message.is_bot:
messages.append(AssistantMessage(content=json.dumps(message.content, default=str)))
messages.append(
AssistantMessage(content=json.dumps(message.content, default=str))
)
else:
messages.append(UserMessage(content=json.dumps(message.content, default=str)))
messages.append(
UserMessage(content=json.dumps(message.content, default=str))
)
# Add the current context
if context_message:
messages.append(SystemMessage(content=context_message))
@@ -156,17 +180,14 @@ async def stream_chat(chat_id: UUID, payload: ChatRequest, db: Session = Depends
messages.append(UserMessage(content=payload.prompt))
async def generate():
response = await client.chat.stream_async(
model=model,
messages=messages
)
response = await client.chat.stream_async(model=model, messages=messages)
async for chunk in response:
if chunk.data.choices[0].delta.content is not None:
content_chunk = chunk.data.choices[0].delta.content
accumulated_content.append(content_chunk)
yield f"data: {json.dumps({'content': content_chunk})}\n\n"
# Create the bot message after all chunks have been processed
chat_message = ChatMessage(
id=uuid4(),
@@ -175,24 +196,25 @@ async def stream_chat(chat_id: UUID, payload: ChatRequest, db: Session = Depends
is_bot=True,
created_at=datetime.utcnow(),
)
db.add(chat_message)
db.commit()
db.refresh(chat_message)
yield "data: [DONE]\n\n"
return StreamingResponse(
generate(),
media_type="text/event-stream"
)
return StreamingResponse(generate(), media_type="text/event-stream")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# Create a new chat
@router.post("/create", response_model=ChatRead, status_code=status.HTTP_201_CREATED)
def create_chat(payload: ChatCreate, db: Session = Depends(get_db), current_user: Profile = Depends(get_current_user)):
def create_chat(
payload: ChatCreate,
db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user),
):
new_chat = Chat(
id=uuid4(),
title=payload.title,
@@ -207,29 +229,42 @@ def create_chat(payload: ChatCreate, db: Session = Depends(get_db), current_user
db.refresh(new_chat)
return new_chat
# Get a chat by ID
@router.get("/{chat_id}", response_model=ChatRead)
def get_chat_by_id(chat_id: UUID, db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user)
):
chat = db.query(Chat).filter(Chat.id == chat_id,
Chat.owner_id == current_user.id
).first()
def get_chat_by_id(
chat_id: UUID,
db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user),
):
chat = (
db.query(Chat)
.filter(Chat.id == chat_id, Chat.owner_id == current_user.id)
.first()
)
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
# Sort messages by created_at in ascending order
chat.messages.sort(key=lambda x: x.created_at)
return chat
# Delete an chat by ID
@router.delete("/{chat_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_chat(chat_id: UUID, db: Session = Depends(get_db), current_user: Profile = Depends(get_current_user)):
chat = db.query(Chat).filter(Chat.id == chat_id, Chat.owner_id == current_user.id).first()
def delete_chat(
chat_id: UUID,
db: Session = Depends(get_db),
current_user: Profile = Depends(get_current_user),
):
chat = (
db.query(Chat)
.filter(Chat.id == chat_id, Chat.owner_id == current_user.id)
.first()
)
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
db.delete(chat)
db.commit()
return None
return None