mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-06-02 11:09:20 +02:00
67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
from uuid import UUID
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from onyx.db.models import AgentSubQuery
|
|
from onyx.db.models import AgentSubQuestion
|
|
|
|
|
|
def create_sub_question(
|
|
db_session: Session,
|
|
chat_session_id: UUID,
|
|
primary_message_id: int,
|
|
sub_question: str,
|
|
sub_answer: str,
|
|
) -> AgentSubQuestion:
|
|
"""Create a new sub-question record in the database."""
|
|
sub_q = AgentSubQuestion(
|
|
chat_session_id=chat_session_id,
|
|
primary_question_id=primary_message_id,
|
|
sub_question=sub_question,
|
|
sub_answer=sub_answer,
|
|
)
|
|
db_session.add(sub_q)
|
|
db_session.flush()
|
|
return sub_q
|
|
|
|
|
|
def create_sub_query(
|
|
db_session: Session,
|
|
chat_session_id: UUID,
|
|
parent_question_id: int,
|
|
sub_query: str,
|
|
) -> AgentSubQuery:
|
|
"""Create a new sub-query record in the database."""
|
|
sub_q = AgentSubQuery(
|
|
chat_session_id=chat_session_id,
|
|
parent_question_id=parent_question_id,
|
|
sub_query=sub_query,
|
|
)
|
|
db_session.add(sub_q)
|
|
db_session.flush()
|
|
return sub_q
|
|
|
|
|
|
def get_sub_questions_for_message(
|
|
db_session: Session,
|
|
primary_message_id: int,
|
|
) -> list[AgentSubQuestion]:
|
|
"""Get all sub-questions for a given primary message."""
|
|
return (
|
|
db_session.query(AgentSubQuestion)
|
|
.filter(AgentSubQuestion.primary_question_id == primary_message_id)
|
|
.all()
|
|
)
|
|
|
|
|
|
def get_sub_queries_for_question(
|
|
db_session: Session,
|
|
sub_question_id: int,
|
|
) -> list[AgentSubQuery]:
|
|
"""Get all sub-queries for a given sub-question."""
|
|
return (
|
|
db_session.query(AgentSubQuery)
|
|
.filter(AgentSubQuery.parent_question_id == sub_question_id)
|
|
.all()
|
|
)
|