reorganization
This commit is contained in:
1
blueprints/__init__.py
Normal file
1
blueprints/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# Blueprints package
|
||||
172
blueprints/conversation/__init__.py
Normal file
172
blueprints/conversation/__init__.py
Normal file
@@ -0,0 +1,172 @@
|
||||
import datetime
|
||||
|
||||
from quart import Blueprint, jsonify, request
|
||||
from quart_jwt_extended import (
|
||||
get_jwt_identity,
|
||||
jwt_refresh_token_required,
|
||||
)
|
||||
|
||||
import blueprints.users.models
|
||||
|
||||
from .agents import main_agent
|
||||
from .logic import (
|
||||
add_message_to_conversation,
|
||||
get_conversation_by_id,
|
||||
rename_conversation,
|
||||
)
|
||||
from .models import (
|
||||
Conversation,
|
||||
PydConversation,
|
||||
PydListConversation,
|
||||
)
|
||||
|
||||
conversation_blueprint = Blueprint(
|
||||
"conversation_api", __name__, url_prefix="/api/conversation"
|
||||
)
|
||||
|
||||
|
||||
@conversation_blueprint.post("/query")
|
||||
@jwt_refresh_token_required
|
||||
async def query():
|
||||
current_user_uuid = get_jwt_identity()
|
||||
user = await blueprints.users.models.User.get(id=current_user_uuid)
|
||||
data = await request.get_json()
|
||||
query = data.get("query")
|
||||
conversation_id = data.get("conversation_id")
|
||||
conversation = await get_conversation_by_id(conversation_id)
|
||||
await conversation.fetch_related("messages")
|
||||
await add_message_to_conversation(
|
||||
conversation=conversation,
|
||||
message=query,
|
||||
speaker="user",
|
||||
user=user,
|
||||
)
|
||||
|
||||
# Build conversation history from recent messages (last 10 for context)
|
||||
recent_messages = (
|
||||
conversation.messages[-10:]
|
||||
if len(conversation.messages) > 10
|
||||
else conversation.messages
|
||||
)
|
||||
|
||||
messages_payload = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": """You are a helpful cat assistant named Simba that understands veterinary terms. When there are questions to you specifically, they are referring to Simba the cat. Answer the user in as if you were a cat named Simba. Don't act too catlike. Be assertive.
|
||||
|
||||
SIMBA FACTS (as of January 2026):
|
||||
- Name: Simba
|
||||
- Species: Feline (Domestic Short Hair / American Short Hair)
|
||||
- Sex: Male, Neutered
|
||||
- Date of Birth: August 8, 2016 (approximately 9 years 5 months old)
|
||||
- Color: Orange
|
||||
- Current Weight: 16 lbs (as of 1/8/2026)
|
||||
- Owner: Ryan Chen
|
||||
- Location: Long Island City, NY
|
||||
- Veterinarian: Court Square Animal Hospital
|
||||
|
||||
Medical Conditions:
|
||||
- Hypertrophic Cardiomyopathy (HCM): Diagnosed 12/11/2025. Concentric left ventricular hypertrophy with no left atrial dilation. Grade II-III/VI systolic heart murmur. No cardiac medications currently needed. Must avoid Domitor, acepromazine, and ketamine during anesthesia.
|
||||
- Dental Issues: Prior extraction of teeth 307 and 407 due to resorption. Tooth 107 extracted on 1/8/2026. Early resorption lesions present on teeth 207, 309, and 409.
|
||||
|
||||
Recent Medical Events:
|
||||
- 1/8/2026: Dental cleaning and tooth 107 extraction. Prescribed Onsior for 3 days. Oravet sealant applied.
|
||||
- 12/11/2025: Echocardiogram confirming HCM diagnosis. Pre-op bloodwork was normal.
|
||||
- 12/1/2025: Visited for decreased appetite/nausea. Received subcutaneous fluids and Cerenia.
|
||||
|
||||
Diet & Lifestyle:
|
||||
- Diet: Hill's I/D wet and dry food
|
||||
- Supplements: Plaque Off
|
||||
- Indoor only cat, only pet in the household
|
||||
|
||||
Upcoming Appointments:
|
||||
- Rabies Vaccine: Due 2/19/2026
|
||||
- Routine Examination: Due 6/1/2026
|
||||
- FVRCP-3yr Vaccine: Due 10/2/2026
|
||||
|
||||
IMPORTANT: When users ask factual questions about Simba's health, medical history, veterinary visits, medications, weight, or any information that would be in documents, you MUST use the simba_search tool to retrieve accurate information before answering. Do not rely on general knowledge - always search the documents for factual questions.""",
|
||||
}
|
||||
]
|
||||
|
||||
# Add recent conversation history
|
||||
for msg in recent_messages[:-1]: # Exclude the message we just added
|
||||
role = "user" if msg.speaker == "user" else "assistant"
|
||||
messages_payload.append({"role": role, "content": msg.text})
|
||||
|
||||
# Add current query
|
||||
messages_payload.append({"role": "user", "content": query})
|
||||
|
||||
payload = {"messages": messages_payload}
|
||||
|
||||
response = await main_agent.ainvoke(payload)
|
||||
message = response.get("messages", [])[-1].content
|
||||
await add_message_to_conversation(
|
||||
conversation=conversation,
|
||||
message=message,
|
||||
speaker="simba",
|
||||
user=user,
|
||||
)
|
||||
return jsonify({"response": message})
|
||||
|
||||
|
||||
@conversation_blueprint.route("/<conversation_id>")
|
||||
@jwt_refresh_token_required
|
||||
async def get_conversation(conversation_id: str):
|
||||
conversation = await Conversation.get(id=conversation_id)
|
||||
current_user_uuid = get_jwt_identity()
|
||||
user = await blueprints.users.models.User.get(id=current_user_uuid)
|
||||
await conversation.fetch_related("messages")
|
||||
|
||||
# Manually serialize the conversation with messages
|
||||
messages = []
|
||||
for msg in conversation.messages:
|
||||
messages.append(
|
||||
{
|
||||
"id": str(msg.id),
|
||||
"text": msg.text,
|
||||
"speaker": msg.speaker.value,
|
||||
"created_at": msg.created_at.isoformat(),
|
||||
}
|
||||
)
|
||||
name = conversation.name
|
||||
if len(messages) > 8 and "datetime" in name.lower():
|
||||
name = await rename_conversation(
|
||||
user=user,
|
||||
conversation=conversation,
|
||||
)
|
||||
print(name)
|
||||
|
||||
return jsonify(
|
||||
{
|
||||
"id": str(conversation.id),
|
||||
"name": name,
|
||||
"messages": messages,
|
||||
"created_at": conversation.created_at.isoformat(),
|
||||
"updated_at": conversation.updated_at.isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@conversation_blueprint.post("/")
|
||||
@jwt_refresh_token_required
|
||||
async def create_conversation():
|
||||
user_uuid = get_jwt_identity()
|
||||
user = await blueprints.users.models.User.get(id=user_uuid)
|
||||
conversation = await Conversation.create(
|
||||
name=f"{user.username} {datetime.datetime.now().timestamp}",
|
||||
user=user,
|
||||
)
|
||||
|
||||
serialized_conversation = await PydConversation.from_tortoise_orm(conversation)
|
||||
return jsonify(serialized_conversation.model_dump())
|
||||
|
||||
|
||||
@conversation_blueprint.get("/")
|
||||
@jwt_refresh_token_required
|
||||
async def get_all_conversations():
|
||||
user_uuid = get_jwt_identity()
|
||||
user = await blueprints.users.models.User.get(id=user_uuid)
|
||||
conversations = Conversation.filter(user=user)
|
||||
serialized_conversations = await PydListConversation.from_queryset(conversations)
|
||||
|
||||
return jsonify(serialized_conversations.model_dump())
|
||||
78
blueprints/conversation/agents.py
Normal file
78
blueprints/conversation/agents.py
Normal file
@@ -0,0 +1,78 @@
|
||||
import os
|
||||
from typing import cast
|
||||
|
||||
from langchain.agents import create_agent
|
||||
from langchain.chat_models import BaseChatModel
|
||||
from langchain.tools import tool
|
||||
from langchain_ollama import ChatOllama
|
||||
from langchain_openai import ChatOpenAI
|
||||
from tavily import AsyncTavilyClient
|
||||
|
||||
from blueprints.rag.logic import query_vector_store
|
||||
|
||||
openai_gpt_5_mini = ChatOpenAI(model="gpt-5-mini")
|
||||
ollama_deepseek = ChatOllama(model="llama3.1:8b", base_url=os.getenv("OLLAMA_URL"))
|
||||
model_with_fallback = cast(
|
||||
BaseChatModel, ollama_deepseek.with_fallbacks([openai_gpt_5_mini])
|
||||
)
|
||||
client = AsyncTavilyClient(os.getenv("TAVILY_KEY"), "")
|
||||
|
||||
|
||||
@tool
|
||||
async def web_search(query: str) -> str:
|
||||
"""Search the web for current information using Tavily.
|
||||
|
||||
Use this tool when you need to:
|
||||
- Find current information not in the knowledge base
|
||||
- Look up recent events, news, or updates
|
||||
- Verify facts or get additional context
|
||||
- Search for information outside of Simba's documents
|
||||
|
||||
Args:
|
||||
query: The search query to look up on the web
|
||||
|
||||
Returns:
|
||||
Search results from the web with titles, content, and source URLs
|
||||
"""
|
||||
response = await client.search(query=query, search_depth="basic")
|
||||
results = response.get("results", [])
|
||||
|
||||
if not results:
|
||||
return "No results found for the query."
|
||||
|
||||
formatted = "\n\n".join(
|
||||
[
|
||||
f"**{result['title']}**\n{result['content']}\nSource: {result['url']}"
|
||||
for result in results[:5]
|
||||
]
|
||||
)
|
||||
return formatted
|
||||
|
||||
|
||||
@tool(response_format="content_and_artifact")
|
||||
async def simba_search(query: str):
|
||||
"""Search through Simba's medical records, veterinary documents, and personal information.
|
||||
|
||||
Use this tool whenever the user asks questions about:
|
||||
- Simba's health history, medical records, or veterinary visits
|
||||
- Medications, treatments, or diagnoses
|
||||
- Weight, diet, or physical characteristics over time
|
||||
- Veterinary recommendations or advice
|
||||
- Ryan's (the owner's) information related to Simba
|
||||
- Any factual information that would be found in documents
|
||||
|
||||
Args:
|
||||
query: The user's question or information need about Simba
|
||||
|
||||
Returns:
|
||||
Relevant information from Simba's documents
|
||||
"""
|
||||
print(f"[SIMBA SEARCH] Tool called with query: {query}")
|
||||
serialized, docs = await query_vector_store(query=query)
|
||||
print(f"[SIMBA SEARCH] Found {len(docs)} documents")
|
||||
print(f"[SIMBA SEARCH] Serialized result length: {len(serialized)}")
|
||||
print(f"[SIMBA SEARCH] First 200 chars: {serialized[:200]}")
|
||||
return serialized, docs
|
||||
|
||||
|
||||
main_agent = create_agent(model=model_with_fallback, tools=[simba_search, web_search])
|
||||
80
blueprints/conversation/logic.py
Normal file
80
blueprints/conversation/logic.py
Normal file
@@ -0,0 +1,80 @@
|
||||
import tortoise.exceptions
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
import blueprints.users.models
|
||||
|
||||
from .models import Conversation, ConversationMessage, RenameConversationOutputSchema
|
||||
|
||||
|
||||
async def create_conversation(name: str = "") -> Conversation:
|
||||
conversation = await Conversation.create(name=name)
|
||||
return conversation
|
||||
|
||||
|
||||
async def add_message_to_conversation(
|
||||
conversation: Conversation,
|
||||
message: str,
|
||||
speaker: str,
|
||||
user: blueprints.users.models.User,
|
||||
) -> ConversationMessage:
|
||||
print(conversation, message, speaker)
|
||||
message = await ConversationMessage.create(
|
||||
text=message,
|
||||
speaker=speaker,
|
||||
conversation=conversation,
|
||||
)
|
||||
|
||||
return message
|
||||
|
||||
|
||||
async def get_the_only_conversation() -> Conversation:
|
||||
try:
|
||||
conversation = await Conversation.all().first()
|
||||
if conversation is None:
|
||||
conversation = await Conversation.create(name="simba_chat")
|
||||
except Exception as _e:
|
||||
conversation = await Conversation.create(name="simba_chat")
|
||||
|
||||
return conversation
|
||||
|
||||
|
||||
async def get_conversation_for_user(user: blueprints.users.models.User) -> Conversation:
|
||||
try:
|
||||
return await Conversation.get(user=user)
|
||||
except tortoise.exceptions.DoesNotExist:
|
||||
await Conversation.get_or_create(name=f"{user.username}'s chat", user=user)
|
||||
|
||||
return await Conversation.get(user=user)
|
||||
|
||||
|
||||
async def get_conversation_by_id(id: str) -> Conversation:
|
||||
return await Conversation.get(id=id)
|
||||
|
||||
|
||||
async def get_conversation_transcript(
|
||||
user: blueprints.users.models.User, conversation: Conversation
|
||||
) -> str:
|
||||
messages = []
|
||||
for message in conversation.messages:
|
||||
messages.append(f"{message.speaker} at {message.created_at}: {message.text}")
|
||||
|
||||
return "\n".join(messages)
|
||||
|
||||
|
||||
async def rename_conversation(
|
||||
user: blueprints.users.models.User,
|
||||
conversation: Conversation,
|
||||
) -> str:
|
||||
messages: str = await get_conversation_transcript(
|
||||
user=user, conversation=conversation
|
||||
)
|
||||
|
||||
llm = ChatOpenAI(model="gpt-4o-mini")
|
||||
structured_llm = llm.with_structured_output(RenameConversationOutputSchema)
|
||||
|
||||
prompt = f"Summarize the following conversation into a sassy one-liner title:\n\n{messages}"
|
||||
response = structured_llm.invoke(prompt)
|
||||
new_name: str = response.get("title", "")
|
||||
conversation.name = new_name
|
||||
await conversation.save()
|
||||
return new_name
|
||||
61
blueprints/conversation/models.py
Normal file
61
blueprints/conversation/models.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import enum
|
||||
from dataclasses import dataclass
|
||||
|
||||
from tortoise import fields
|
||||
from tortoise.contrib.pydantic import (
|
||||
pydantic_model_creator,
|
||||
pydantic_queryset_creator,
|
||||
)
|
||||
from tortoise.models import Model
|
||||
|
||||
|
||||
@dataclass
|
||||
class RenameConversationOutputSchema:
|
||||
title: str
|
||||
justification: str
|
||||
|
||||
|
||||
class Speaker(enum.Enum):
|
||||
USER = "user"
|
||||
SIMBA = "simba"
|
||||
|
||||
|
||||
class Conversation(Model):
|
||||
id = fields.UUIDField(primary_key=True)
|
||||
name = fields.CharField(max_length=255)
|
||||
created_at = fields.DatetimeField(auto_now_add=True)
|
||||
updated_at = fields.DatetimeField(auto_now=True)
|
||||
user: fields.ForeignKeyRelation = fields.ForeignKeyField(
|
||||
"models.User", related_name="conversations", null=True
|
||||
)
|
||||
|
||||
class Meta:
|
||||
table = "conversations"
|
||||
|
||||
|
||||
class ConversationMessage(Model):
|
||||
id = fields.UUIDField(primary_key=True)
|
||||
text = fields.TextField()
|
||||
conversation = fields.ForeignKeyField(
|
||||
"models.Conversation", related_name="messages"
|
||||
)
|
||||
created_at = fields.DatetimeField(auto_now_add=True)
|
||||
speaker = fields.CharEnumField(enum_type=Speaker, max_length=10)
|
||||
|
||||
class Meta:
|
||||
table = "conversation_messages"
|
||||
|
||||
|
||||
PydConversationMessage = pydantic_model_creator(ConversationMessage)
|
||||
PydConversation = pydantic_model_creator(
|
||||
Conversation, name="Conversation", allow_cycles=True, exclude=("user",)
|
||||
)
|
||||
PydConversationWithMessages = pydantic_model_creator(
|
||||
Conversation,
|
||||
name="ConversationWithMessages",
|
||||
allow_cycles=True,
|
||||
exclude=("user",),
|
||||
include=("messages",),
|
||||
)
|
||||
PydListConversation = pydantic_queryset_creator(Conversation)
|
||||
PydListConversationMessage = pydantic_queryset_creator(ConversationMessage)
|
||||
47
blueprints/rag/__init__.py
Normal file
47
blueprints/rag/__init__.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from quart import Blueprint, jsonify
|
||||
from quart_jwt_extended import jwt_refresh_token_required
|
||||
|
||||
from .logic import get_vector_store_stats, index_documents, vector_store
|
||||
from blueprints.users.decorators import admin_required
|
||||
|
||||
rag_blueprint = Blueprint("rag_api", __name__, url_prefix="/api/rag")
|
||||
|
||||
|
||||
@rag_blueprint.get("/stats")
|
||||
@jwt_refresh_token_required
|
||||
async def get_stats():
|
||||
"""Get vector store statistics."""
|
||||
stats = get_vector_store_stats()
|
||||
return jsonify(stats)
|
||||
|
||||
|
||||
@rag_blueprint.post("/index")
|
||||
@admin_required
|
||||
async def trigger_index():
|
||||
"""Trigger indexing of documents from Paperless-NGX. Admin only."""
|
||||
try:
|
||||
await index_documents()
|
||||
stats = get_vector_store_stats()
|
||||
return jsonify({"status": "success", "stats": stats})
|
||||
except Exception as e:
|
||||
return jsonify({"status": "error", "message": str(e)}), 500
|
||||
|
||||
|
||||
@rag_blueprint.post("/reindex")
|
||||
@admin_required
|
||||
async def trigger_reindex():
|
||||
"""Clear and reindex all documents. Admin only."""
|
||||
try:
|
||||
# Clear existing documents
|
||||
collection = vector_store._collection
|
||||
all_docs = collection.get()
|
||||
|
||||
if all_docs["ids"]:
|
||||
collection.delete(ids=all_docs["ids"])
|
||||
|
||||
# Reindex
|
||||
await index_documents()
|
||||
stats = get_vector_store_stats()
|
||||
return jsonify({"status": "success", "stats": stats})
|
||||
except Exception as e:
|
||||
return jsonify({"status": "error", "message": str(e)}), 500
|
||||
75
blueprints/rag/fetchers.py
Normal file
75
blueprints/rag/fetchers.py
Normal file
@@ -0,0 +1,75 @@
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class PaperlessNGXService:
|
||||
def __init__(self):
|
||||
self.base_url = os.getenv("BASE_URL")
|
||||
self.token = os.getenv("PAPERLESS_TOKEN")
|
||||
self.url = f"http://{os.getenv('BASE_URL')}/api/documents/?tags__id=8"
|
||||
self.headers = {"Authorization": f"Token {os.getenv('PAPERLESS_TOKEN')}"}
|
||||
|
||||
def get_data(self):
|
||||
print(f"Getting data from: {self.url}")
|
||||
r = httpx.get(self.url, headers=self.headers)
|
||||
results = r.json()["results"]
|
||||
|
||||
nextLink = r.json().get("next")
|
||||
|
||||
while nextLink:
|
||||
r = httpx.get(nextLink, headers=self.headers)
|
||||
results += r.json()["results"]
|
||||
nextLink = r.json().get("next")
|
||||
|
||||
return results
|
||||
|
||||
def get_doc_by_id(self, doc_id: int):
|
||||
url = f"http://{os.getenv('BASE_URL')}/api/documents/{doc_id}/"
|
||||
r = httpx.get(url, headers=self.headers)
|
||||
return r.json()
|
||||
|
||||
def download_pdf_from_id(self, id: int) -> str:
|
||||
download_url = f"http://{os.getenv('BASE_URL')}/api/documents/{id}/download/"
|
||||
response = httpx.get(
|
||||
download_url, headers=self.headers, follow_redirects=True, timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
# Use a temporary file for the downloaded PDF
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
|
||||
temp_file.write(response.content)
|
||||
temp_file.close()
|
||||
temp_pdf_path = temp_file.name
|
||||
pdf_to_process = temp_pdf_path
|
||||
return pdf_to_process
|
||||
|
||||
def upload_cleaned_content(self, document_id, data):
|
||||
PUTS_URL = f"http://{os.getenv('BASE_URL')}/api/documents/{document_id}/"
|
||||
r = httpx.put(PUTS_URL, headers=self.headers, data=data)
|
||||
r.raise_for_status()
|
||||
|
||||
def upload_description(self, description_filepath, file, title, exif_date: str):
|
||||
POST_URL = f"http://{os.getenv('BASE_URL')}/api/documents/post_document/"
|
||||
files = {"document": ("description_filepath", file, "application/txt")}
|
||||
data = {
|
||||
"title": title,
|
||||
"create": exif_date,
|
||||
"document_type": 3,
|
||||
"tags": [7],
|
||||
}
|
||||
|
||||
r = httpx.post(POST_URL, headers=self.headers, data=data, files=files)
|
||||
r.raise_for_status()
|
||||
|
||||
def get_tags(self):
|
||||
GET_URL = f"http://{os.getenv('BASE_URL')}/api/tags/"
|
||||
r = httpx.get(GET_URL, headers=self.headers)
|
||||
data = r.json()
|
||||
return {tag["id"]: tag["name"] for tag in data["results"]}
|
||||
|
||||
def get_doctypes(self):
|
||||
GET_URL = f"http://{os.getenv('BASE_URL')}/api/document_types/"
|
||||
r = httpx.get(GET_URL, headers=self.headers)
|
||||
data = r.json()
|
||||
return {doctype["id"]: doctype["name"] for doctype in data["results"]}
|
||||
101
blueprints/rag/logic.py
Normal file
101
blueprints/rag/logic.py
Normal file
@@ -0,0 +1,101 @@
|
||||
import datetime
|
||||
import os
|
||||
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_core.documents import Document
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
from .fetchers import PaperlessNGXService
|
||||
|
||||
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
|
||||
|
||||
vector_store = Chroma(
|
||||
collection_name="simba_docs",
|
||||
embedding_function=embeddings,
|
||||
persist_directory=os.getenv("CHROMADB_PATH", ""),
|
||||
)
|
||||
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000, # chunk size (characters)
|
||||
chunk_overlap=200, # chunk overlap (characters)
|
||||
add_start_index=True, # track index in original document
|
||||
)
|
||||
|
||||
|
||||
def date_to_epoch(date_str: str) -> float:
|
||||
split_date = date_str.split("-")
|
||||
date = datetime.datetime(
|
||||
int(split_date[0]),
|
||||
int(split_date[1]),
|
||||
int(split_date[2]),
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
)
|
||||
|
||||
return date.timestamp()
|
||||
|
||||
|
||||
async def fetch_documents_from_paperless_ngx() -> list[Document]:
|
||||
ppngx = PaperlessNGXService()
|
||||
data = ppngx.get_data()
|
||||
doctypes = ppngx.get_doctypes()
|
||||
documents = []
|
||||
for doc in data:
|
||||
metadata = {
|
||||
"created_date": date_to_epoch(doc["created_date"]),
|
||||
"filename": doc["original_file_name"],
|
||||
"document_type": doctypes.get(doc["document_type"], ""),
|
||||
}
|
||||
documents.append(Document(page_content=doc["content"], metadata=metadata))
|
||||
|
||||
return documents
|
||||
|
||||
|
||||
async def index_documents():
|
||||
documents = await fetch_documents_from_paperless_ngx()
|
||||
|
||||
splits = text_splitter.split_documents(documents)
|
||||
await vector_store.aadd_documents(documents=splits)
|
||||
|
||||
|
||||
async def query_vector_store(query: str):
|
||||
retrieved_docs = await vector_store.asimilarity_search(query, k=2)
|
||||
serialized = "\n\n".join(
|
||||
(f"Source: {doc.metadata}\nContent: {doc.page_content}")
|
||||
for doc in retrieved_docs
|
||||
)
|
||||
return serialized, retrieved_docs
|
||||
|
||||
|
||||
def get_vector_store_stats():
|
||||
"""Get statistics about the vector store."""
|
||||
collection = vector_store._collection
|
||||
count = collection.count()
|
||||
return {
|
||||
"total_documents": count,
|
||||
"collection_name": collection.name,
|
||||
}
|
||||
|
||||
|
||||
def list_all_documents(limit: int = 10):
|
||||
"""List documents in the vector store with their metadata."""
|
||||
collection = vector_store._collection
|
||||
results = collection.get(limit=limit, include=["metadatas", "documents"])
|
||||
|
||||
documents = []
|
||||
for i, doc_id in enumerate(results["ids"]):
|
||||
documents.append(
|
||||
{
|
||||
"id": doc_id,
|
||||
"metadata": results["metadatas"][i]
|
||||
if results.get("metadatas")
|
||||
else None,
|
||||
"content_preview": results["documents"][i][:200]
|
||||
if results.get("documents")
|
||||
else None,
|
||||
}
|
||||
)
|
||||
|
||||
return documents
|
||||
0
blueprints/rag/models.py
Normal file
0
blueprints/rag/models.py
Normal file
188
blueprints/users/__init__.py
Normal file
188
blueprints/users/__init__.py
Normal file
@@ -0,0 +1,188 @@
|
||||
from quart import Blueprint, jsonify, request
|
||||
from quart_jwt_extended import (
|
||||
create_access_token,
|
||||
create_refresh_token,
|
||||
jwt_refresh_token_required,
|
||||
get_jwt_identity,
|
||||
)
|
||||
from .models import User
|
||||
from .oidc_service import OIDCUserService
|
||||
from config.oidc_config import oidc_config
|
||||
import secrets
|
||||
import httpx
|
||||
from urllib.parse import urlencode
|
||||
import hashlib
|
||||
import base64
|
||||
|
||||
|
||||
user_blueprint = Blueprint("user_api", __name__, url_prefix="/api/user")
|
||||
|
||||
# In-memory storage for OIDC state/PKCE (production: use Redis or database)
|
||||
# Format: {state: {"pkce_verifier": str, "redirect_after_login": str}}
|
||||
_oidc_sessions = {}
|
||||
|
||||
|
||||
@user_blueprint.route("/oidc/login", methods=["GET"])
|
||||
async def oidc_login():
|
||||
"""
|
||||
Initiate OIDC login flow
|
||||
Generates PKCE parameters and redirects to Authelia
|
||||
"""
|
||||
if not oidc_config.validate_config():
|
||||
return jsonify({"error": "OIDC not configured"}), 500
|
||||
|
||||
try:
|
||||
# Generate PKCE parameters
|
||||
code_verifier = secrets.token_urlsafe(64)
|
||||
|
||||
# For PKCE, we need code_challenge = BASE64URL(SHA256(code_verifier))
|
||||
code_challenge = (
|
||||
base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest())
|
||||
.decode()
|
||||
.rstrip("=")
|
||||
)
|
||||
|
||||
# Generate state for CSRF protection
|
||||
state = secrets.token_urlsafe(32)
|
||||
|
||||
# Store PKCE verifier and state for callback validation
|
||||
_oidc_sessions[state] = {
|
||||
"pkce_verifier": code_verifier,
|
||||
"redirect_after_login": request.args.get("redirect", "/"),
|
||||
}
|
||||
|
||||
# Get authorization endpoint from discovery
|
||||
discovery = await oidc_config.get_discovery_document()
|
||||
auth_endpoint = discovery.get("authorization_endpoint")
|
||||
|
||||
# Build authorization URL
|
||||
params = {
|
||||
"client_id": oidc_config.client_id,
|
||||
"response_type": "code",
|
||||
"redirect_uri": oidc_config.redirect_uri,
|
||||
"scope": "openid email profile groups",
|
||||
"state": state,
|
||||
"code_challenge": code_challenge,
|
||||
"code_challenge_method": "S256",
|
||||
}
|
||||
|
||||
auth_url = f"{auth_endpoint}?{urlencode(params)}"
|
||||
|
||||
return jsonify({"auth_url": auth_url})
|
||||
except Exception as e:
|
||||
return jsonify({"error": f"OIDC login failed: {str(e)}"}), 500
|
||||
|
||||
|
||||
@user_blueprint.route("/oidc/callback", methods=["GET"])
|
||||
async def oidc_callback():
|
||||
"""
|
||||
Handle OIDC callback from Authelia
|
||||
Exchanges authorization code for tokens, verifies ID token, and creates/updates user
|
||||
"""
|
||||
# Get authorization code and state from callback
|
||||
code = request.args.get("code")
|
||||
state = request.args.get("state")
|
||||
error = request.args.get("error")
|
||||
|
||||
if error:
|
||||
return jsonify({"error": f"OIDC error: {error}"}), 400
|
||||
|
||||
if not code or not state:
|
||||
return jsonify({"error": "Missing code or state"}), 400
|
||||
|
||||
# Validate state and retrieve PKCE verifier
|
||||
session = _oidc_sessions.pop(state, None)
|
||||
if not session:
|
||||
return jsonify({"error": "Invalid or expired state"}), 400
|
||||
|
||||
pkce_verifier = session["pkce_verifier"]
|
||||
|
||||
# Exchange authorization code for tokens
|
||||
discovery = await oidc_config.get_discovery_document()
|
||||
token_endpoint = discovery.get("token_endpoint")
|
||||
|
||||
token_data = {
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": oidc_config.redirect_uri,
|
||||
"client_id": oidc_config.client_id,
|
||||
"client_secret": oidc_config.client_secret,
|
||||
"code_verifier": pkce_verifier,
|
||||
}
|
||||
|
||||
# Use client_secret_post method (credentials in POST body)
|
||||
async with httpx.AsyncClient() as client:
|
||||
token_response = await client.post(token_endpoint, data=token_data)
|
||||
|
||||
if token_response.status_code != 200:
|
||||
return jsonify(
|
||||
{"error": f"Failed to exchange code for token: {token_response.text}"}
|
||||
), 400
|
||||
|
||||
tokens = token_response.json()
|
||||
|
||||
id_token = tokens.get("id_token")
|
||||
if not id_token:
|
||||
return jsonify({"error": "No ID token received"}), 400
|
||||
|
||||
# Verify ID token
|
||||
try:
|
||||
claims = await oidc_config.verify_id_token(id_token)
|
||||
except Exception as e:
|
||||
return jsonify({"error": f"ID token verification failed: {str(e)}"}), 400
|
||||
|
||||
# Get or create user from OIDC claims
|
||||
user = await OIDCUserService.get_or_create_user_from_oidc(claims)
|
||||
|
||||
# Issue backend JWT tokens
|
||||
access_token = create_access_token(identity=str(user.id))
|
||||
refresh_token = create_refresh_token(identity=str(user.id))
|
||||
|
||||
# Return tokens to frontend
|
||||
# Frontend will handle storing these and redirecting
|
||||
return jsonify(
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
user={
|
||||
"id": str(user.id),
|
||||
"username": user.username,
|
||||
"email": user.email,
|
||||
"groups": user.ldap_groups,
|
||||
"is_admin": user.is_admin(),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@user_blueprint.route("/refresh", methods=["POST"])
|
||||
@jwt_refresh_token_required
|
||||
async def refresh():
|
||||
"""Refresh access token (unchanged from original)"""
|
||||
user_id = get_jwt_identity()
|
||||
new_token = create_access_token(identity=user_id)
|
||||
return jsonify(access_token=new_token)
|
||||
|
||||
|
||||
# Legacy username/password login - kept for backward compatibility during migration
|
||||
@user_blueprint.route("/login", methods=["POST"])
|
||||
async def login():
|
||||
"""
|
||||
Legacy username/password login
|
||||
This can be removed after full OIDC migration is complete
|
||||
"""
|
||||
data = await request.get_json()
|
||||
username = data.get("username")
|
||||
password = data.get("password")
|
||||
|
||||
user = await User.filter(username=username).first()
|
||||
|
||||
if not user or not user.verify_password(password):
|
||||
return jsonify({"msg": "Invalid credentials"}), 401
|
||||
|
||||
access_token = create_access_token(identity=str(user.id))
|
||||
refresh_token = create_refresh_token(identity=str(user.id))
|
||||
|
||||
return jsonify(
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
user={"id": str(user.id), "username": user.username},
|
||||
)
|
||||
26
blueprints/users/decorators.py
Normal file
26
blueprints/users/decorators.py
Normal file
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
Authentication decorators for role-based access control.
|
||||
"""
|
||||
|
||||
from functools import wraps
|
||||
from quart import jsonify
|
||||
from quart_jwt_extended import jwt_refresh_token_required, get_jwt_identity
|
||||
from .models import User
|
||||
|
||||
|
||||
def admin_required(fn):
|
||||
"""
|
||||
Decorator that requires the user to be an admin (member of lldap_admin group).
|
||||
Must be used on async route handlers.
|
||||
"""
|
||||
|
||||
@wraps(fn)
|
||||
@jwt_refresh_token_required
|
||||
async def wrapper(*args, **kwargs):
|
||||
user_id = get_jwt_identity()
|
||||
user = await User.get_or_none(id=user_id)
|
||||
if not user or not user.is_admin():
|
||||
return jsonify({"error": "Admin access required"}), 403
|
||||
return await fn(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
46
blueprints/users/models.py
Normal file
46
blueprints/users/models.py
Normal file
@@ -0,0 +1,46 @@
|
||||
from tortoise.models import Model
|
||||
from tortoise import fields
|
||||
|
||||
|
||||
import bcrypt
|
||||
|
||||
|
||||
class User(Model):
|
||||
id = fields.UUIDField(primary_key=True)
|
||||
username = fields.CharField(max_length=255)
|
||||
password = fields.BinaryField(null=True) # Hashed - nullable for OIDC users
|
||||
email = fields.CharField(max_length=100, unique=True)
|
||||
|
||||
# OIDC fields
|
||||
oidc_subject = fields.CharField(
|
||||
max_length=255, unique=True, null=True, index=True
|
||||
) # "sub" claim from OIDC
|
||||
auth_provider = fields.CharField(
|
||||
max_length=50, default="local"
|
||||
) # "local" or "oidc"
|
||||
ldap_groups = fields.JSONField(default=[]) # LDAP groups from OIDC claims
|
||||
|
||||
created_at = fields.DatetimeField(auto_now_add=True)
|
||||
updated_at = fields.DatetimeField(auto_now=True)
|
||||
|
||||
class Meta:
|
||||
table = "users"
|
||||
|
||||
def has_group(self, group: str) -> bool:
|
||||
"""Check if user belongs to a specific LDAP group."""
|
||||
return group in (self.ldap_groups or [])
|
||||
|
||||
def is_admin(self) -> bool:
|
||||
"""Check if user is an admin (member of lldap_admin group)."""
|
||||
return self.has_group("lldap_admin")
|
||||
|
||||
def set_password(self, plain_password: str):
|
||||
self.password = bcrypt.hashpw(
|
||||
plain_password.encode("utf-8"),
|
||||
bcrypt.gensalt(),
|
||||
)
|
||||
|
||||
def verify_password(self, plain_password: str):
|
||||
if not self.password:
|
||||
return False
|
||||
return bcrypt.checkpw(plain_password.encode("utf-8"), self.password)
|
||||
81
blueprints/users/oidc_service.py
Normal file
81
blueprints/users/oidc_service.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
OIDC User Management Service
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional
|
||||
from uuid import uuid4
|
||||
from .models import User
|
||||
|
||||
|
||||
class OIDCUserService:
|
||||
"""Service for managing OIDC user authentication and provisioning"""
|
||||
|
||||
@staticmethod
|
||||
async def get_or_create_user_from_oidc(claims: Dict[str, Any]) -> User:
|
||||
"""
|
||||
Get existing user by OIDC subject, or create new user from OIDC claims
|
||||
|
||||
Args:
|
||||
claims: Decoded OIDC ID token claims
|
||||
|
||||
Returns:
|
||||
User object (existing or newly created)
|
||||
"""
|
||||
oidc_subject = claims.get("sub")
|
||||
if not oidc_subject:
|
||||
raise ValueError("No 'sub' claim in ID token")
|
||||
|
||||
# Try to find existing user by OIDC subject
|
||||
user = await User.filter(oidc_subject=oidc_subject).first()
|
||||
|
||||
if user:
|
||||
# Update user info from latest claims (optional)
|
||||
user.email = claims.get("email", user.email)
|
||||
user.username = (
|
||||
claims.get("preferred_username") or claims.get("name") or user.username
|
||||
)
|
||||
# Update LDAP groups from claims
|
||||
user.ldap_groups = claims.get("groups", [])
|
||||
await user.save()
|
||||
return user
|
||||
|
||||
# Check if user exists by email (migration case)
|
||||
email = claims.get("email")
|
||||
if email:
|
||||
user = await User.filter(email=email, auth_provider="local").first()
|
||||
if user:
|
||||
# Migrate existing local user to OIDC
|
||||
user.oidc_subject = oidc_subject
|
||||
user.auth_provider = "oidc"
|
||||
user.password = None # Clear password
|
||||
user.ldap_groups = claims.get("groups", [])
|
||||
await user.save()
|
||||
return user
|
||||
|
||||
# Create new user from OIDC claims
|
||||
username = (
|
||||
claims.get("preferred_username")
|
||||
or claims.get("name")
|
||||
or claims.get("email", "").split("@")[0]
|
||||
or f"user_{oidc_subject[:8]}"
|
||||
)
|
||||
|
||||
# Extract LDAP groups from claims
|
||||
groups = claims.get("groups", [])
|
||||
|
||||
user = await User.create(
|
||||
id=uuid4(),
|
||||
username=username,
|
||||
email=email or f"{oidc_subject}@oidc.local", # Fallback if no email claim
|
||||
oidc_subject=oidc_subject,
|
||||
auth_provider="oidc",
|
||||
password=None,
|
||||
ldap_groups=groups,
|
||||
)
|
||||
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
async def find_user_by_oidc_subject(oidc_subject: str) -> Optional[User]:
|
||||
"""Find user by OIDC subject ID"""
|
||||
return await User.filter(oidc_subject=oidc_subject).first()
|
||||
Reference in New Issue
Block a user