vynte-infra/plane
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[PAI-930] feat: rewrite pi response (#4627)

Browse Source
This commit is contained in:
Akhil Vamshi Konam
2025-11-10 13:53:58 +05:30
committed by GitHub
parent aa1fd6188f
commit 8a147dbfd6
12 changed files with 582 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/pi/pi/alembic/env.py
+1
View File
@@ -19,6 +19,7 @@ from pi.app.models import WorkspaceVectorization # noqa: F401
from pi.app.models.action_artifact import ActionArtifact # noqa: F401
from pi.app.models.action_artifact import ActionArtifactVersion # noqa: F401
from pi.app.models.base import BaseModel
from pi.app.models.message_clarification import MessageClarification # noqa: F401
from pi.app.models.oauth import PlaneOAuthState # noqa: F401
from pi.app.models.oauth import PlaneOAuthToken # noqa: F401
apps/pi/pi/alembic/versions/6082d91ecc4e_rewrite.py
+23
View File
@@ -0,0 +1,23 @@
from typing import Sequence
from typing import Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "6082d91ecc4e"
down_revision: Union[str, None] = "c71d3b70bf59"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.add_column("messages", sa.Column("is_replaced", sa.Boolean(), server_default=sa.text("false"), nullable=False))
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column("messages", "is_replaced")
# ### end Alembic commands ###
apps/pi/pi/app/api/v1/endpoints/chat.py
+83 -41
View File
@@ -73,6 +73,9 @@ from pi.services.retrievers.pg_store import retrieve_chat_history
from pi.services.retrievers.pg_store import soft_delete_chat
from pi.services.retrievers.pg_store import unfavorite_chat
from pi.services.retrievers.pg_store import update_message_feedback
from pi.services.retrievers.pg_store.message import get_message_by_id
from pi.services.retrievers.pg_store.message import mark_assistant_response_as_replaced
from pi.services.retrievers.pg_store.message import reconstruct_chat_request_from_message
from pi.services.retrievers.pg_store.message import upsert_message
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps
@@ -819,7 +822,10 @@ async def stream_answer(token: UUID4, db: AsyncSession = Depends(get_async_sessi
"""Second phase of two-step flow.
Looks up the queued ChatRequest by token (message_id), deletes the queue
entry, and then re-uses the existing /get-answer/ logic to start the SSE
stream via a pure GET endpoint."""
stream via a pure GET endpoint.
Also handles REGENERATE when called with an existing user message ID
(no QUEUE flow step present)."""
try:
auth = await is_valid_session(session)
if not auth.user:
@@ -838,53 +844,89 @@ async def stream_answer(token: UUID4, db: AsyncSession = Depends(get_async_sessi
flow_step: MessageFlowStep | None = res.scalar_one_or_none()
if not flow_step:
return JSONResponse(status_code=404, content={"detail": "Unknown or expired stream token"})
# REGENERATE FLOW: No QUEUE flow step means this is regenerate
log.info(f"No QUEUE flow step found for token {token}. Treating as regenerate request.")
# 1. Get the user message
user_message = await get_message_by_id(db, token)
# Parse the stored ChatRequest
try:
raw_data = flow_step.execution_data or {}
# Check if this is an OAuth message (has oauth_required field)
if flow_step.oauth_required:
# This is an OAuth message - check if OAuth is now complete
# by looking at the oauth_completed column
if flow_step.oauth_completed:
# OAuth is now complete! Process the request normally
log.info(f"OAuth completed for message {token}. Processing request.")
if not user_message:
log.warning(f"Message not found for token {token}")
return JSONResponse(status_code=404, content={"detail": "Message not found"})
# Control continues to the normal processing code below
else:
# OAuth is still required
return JSONResponse(
status_code=401,
content={
"detail": "OAuth authorization still required. Please complete OAuth authentication first.",
"error_code": "OAUTH_REQUIRED",
},
)
# 2. Mark old assistant response as replaced BEFORE generating new one
# This ensures retrieve_chat_history() won't include it in context for new generation
marked = await mark_assistant_response_as_replaced(db, user_message.id)
if marked:
log.info(f"Marked old assistant response(s) as replaced for message {token}")
else:
log.info(f"No existing assistant response found for message {token} (first generation or already replaced)")
# Convert empty-string UUIDs to None so pydantic validation passes
for field in ["project_id", "workspace_id", "chat_id"]:
if field in raw_data and raw_data[field] == "":
raw_data[field] = None
raw_data["user_id"] = user_id
queued_request = ChatRequest.parse_obj(raw_data)
# 4. Reconstruct ChatRequest from user message
try:
queued_request = await reconstruct_chat_request_from_message(db, user_message, user_id)
except Exception as e:
log.error(f"Error reconstructing ChatRequest from message {token}: {e}")
return JSONResponse(status_code=500, content={"detail": "Failed to reconstruct request"})
except Exception as e:
log.error(f"Malformed execution_data for token {token}: {e!s}")
return JSONResponse(status_code=500, content={"detail": "Corrupted queued request"})
# 5. Pass token_id so new assistant message reuses same user message as parent
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Failed to attach token_id to regenerate request context: {e!s}")
# Consume the queue entry so the token is single-use
await db.delete(flow_step)
await db.commit()
log.info(f"Regenerating response for message {token}")
# Pass the token/message_id forward so downstream processing reuses the same Message row
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Failed to attach token_id to queued request context: {e!s}")
# 6. Stream new response (get_answer will call process_query_stream)
# When process_query_stream calls retrieve_chat_history,
# it will NOT see the old response because is_replaced=True
return await get_answer(data=queued_request, session=session)
# Delegate to existing get_answer for streaming
return await get_answer(data=queued_request, session=session)
else:
# NORMAL FLOW: QUEUE flow step exists, this is first generation
# Parse the stored ChatRequest
try:
raw_data = flow_step.execution_data or {}
# Check if this is an OAuth message (has oauth_required field)
if flow_step.oauth_required:
# This is an OAuth message - check if OAuth is now complete
# by looking at the oauth_completed column
if flow_step.oauth_completed:
# OAuth is now complete! Process the request normally
log.info(f"OAuth completed for message {token}. Processing request.")
# Control continues to the normal processing code below
else:
# OAuth is still required
return JSONResponse(
status_code=401,
content={
"detail": "OAuth authorization still required. Please complete OAuth authentication first.",
"error_code": "OAUTH_REQUIRED",
},
)
# Convert empty-string UUIDs to None so pydantic validation passes
for field in ["project_id", "workspace_id", "chat_id"]:
if field in raw_data and raw_data[field] == "":
raw_data[field] = None
raw_data["user_id"] = user_id
queued_request = ChatRequest.parse_obj(raw_data)
except Exception as e:
log.error(f"Malformed execution_data for token {token}: {e!s}")
return JSONResponse(status_code=500, content={"detail": "Corrupted queued request"})
# Consume the queue entry so the token is single-use (for normal flow)
await db.delete(flow_step)
await db.commit()
# Pass the token/message_id forward so downstream processing reuses the same Message row
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Failed to attach token_id to queued request context: {e!s}")
# Delegate to existing get_answer for streaming
return await get_answer(data=queued_request, session=session)
@router.post("/execute-action/")
apps/pi/pi/app/api/v1/endpoints/mobile/chat.py
+293
View File
@@ -18,7 +18,12 @@ from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.api.v1.dependencies import jwt_schema
from pi.app.api.v1.dependencies import validate_jwt_token
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import UserTypeChoices
from pi.app.models.message import MessageFlowStep
from pi.app.schemas.mobile.chat import ChatFeedbackMobile
from pi.app.schemas.mobile.chat import ChatInitializationRequestMobile
from pi.app.schemas.mobile.chat import ChatRequestMobile
from pi.app.schemas.mobile.chat import ChatSearchResponseMobile
from pi.app.schemas.mobile.chat import ChatStartResponseMobile
@@ -31,6 +36,8 @@ from pi.app.schemas.mobile.chat import ModelsResponseMobile
from pi.app.schemas.mobile.chat import RenameChatRequestMobile
from pi.app.schemas.mobile.chat import TitleRequestMobile
from pi.app.schemas.mobile.chat import UnfavoriteChatRequestMobile
from pi.app.utils import validate_chat_initialization
from pi.app.utils import validate_chat_request
from pi.app.utils.background_tasks import schedule_chat_deletion
from pi.app.utils.background_tasks import schedule_chat_rename
from pi.app.utils.background_tasks import schedule_chat_search_upsert
@@ -41,6 +48,8 @@ from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.tool_utils import format_clarification_as_text
from pi.services.chat.search import ChatSearchService
from pi.services.chat.templates import tiles_factory
from pi.services.chat.utils import initialize_new_chat
from pi.services.chat.utils import resolve_workspace_slug
from pi.services.retrievers.pg_store import favorite_chat
from pi.services.retrievers.pg_store import get_active_models
from pi.services.retrievers.pg_store import get_chat_messages
@@ -51,6 +60,11 @@ from pi.services.retrievers.pg_store import retrieve_chat_history
from pi.services.retrievers.pg_store import soft_delete_chat
from pi.services.retrievers.pg_store import unfavorite_chat
from pi.services.retrievers.pg_store import update_message_feedback
from pi.services.retrievers.pg_store.message import get_message_by_id
from pi.services.retrievers.pg_store.message import mark_assistant_response_as_replaced
from pi.services.retrievers.pg_store.message import reconstruct_chat_request_from_message
from pi.services.retrievers.pg_store.message import upsert_message
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps
log = logger.getChild("v1/mobile/chat")
mobile_router = APIRouter()
@@ -183,6 +197,285 @@ async def get_answer(
return JSONResponse(status_code=500, content={"detail": "Internal Server Error"})
@mobile_router.post("/initialize-chat/")
async def initialize_chat(
data: ChatInitializationRequestMobile,
db: AsyncSession = Depends(get_async_session),
token: HTTPAuthorizationCredentials = Depends(jwt_schema),
):
"""Initialize a new chat and return the chat_id immediately."""
try:
auth = await validate_jwt_token(token)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
user_id = auth.user.id
except Exception as e:
log.error(f"Error validating JWT: {e!s}")
return JSONResponse(status_code=401, content={"detail": "Invalid JWT"})
# Validate request data
validation_error = validate_chat_initialization(data)
if validation_error:
return JSONResponse(status_code=validation_error["status_code"], content={"detail": validation_error["detail"]})
# Initialize chat using standalone function (workspace details backfilled later in queue_answer)
result = await initialize_new_chat(
user_id=user_id,
db=db,
chat_id=data.chat_id,
is_project_chat=data.is_project_chat,
workspace_in_context=data.workspace_in_context,
workspace_id=data.workspace_id,
)
# Handle result from service layer
if result["success"]:
return JSONResponse(content={"chat_id": result["chat_id"]})
else:
# Map service error codes to HTTP status codes
status_code = 500 # default
if result.get("error_code") == "CHAT_EXISTS":
status_code = 409
elif result.get("error_code") == "CHAT_CREATION_FAILED":
status_code = 500
elif result.get("error_code") == "UNEXPECTED_ERROR":
status_code = 500
return JSONResponse(status_code=status_code, content={"detail": result["message"]})
@mobile_router.post("/queue-answer/")
async def queue_answer(
data: ChatRequestMobile,
db: AsyncSession = Depends(get_async_session),
token: HTTPAuthorizationCredentials = Depends(jwt_schema),
):
"""First phase of two-step streaming flow.
Persists the ChatRequest payload and returns a one-time stream token.
The token is simply the UUID of a freshly created *user* Message row so we
don't need a new table. The rest of the ChatRequest fields are stored in
a MessageFlowStep (tool_name="QUEUE", step_order=0) until the client
later redeems the token via /stream-answer/{token}."""
try:
auth = await validate_jwt_token(token)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
except Exception as e:
log.error(f"Error validating JWT: {e!s}")
return JSONResponse(status_code=401, content={"detail": "Invalid JWT"})
validation_error = validate_chat_request(data)
if validation_error:
return JSONResponse(status_code=validation_error["status_code"], content={"detail": validation_error["detail"]})
# 1. Create a new USER message that will serve as the token
if data.chat_id is None:
return JSONResponse(status_code=400, content={"detail": "chat_id is required. Call /initialize-chat/ first."})
## need to resolve ws id and slug from project id if it's project chat
if not data.workspace_id:
if data.is_project_chat and data.project_id:
log.info(f"Mobile Queue-answer: Resolving workspace_id from project_id: {data.project_id}")
resolved_workspace_id = await resolve_workspace_id_from_project_id(str(data.project_id))
# The DB may return an asyncpg UUID object. Convert safely to a standard uuid.UUID.
from uuid import UUID
workspace_id_to_use = UUID(str(resolved_workspace_id)) if resolved_workspace_id else None
log.info(f"Mobile Queue-answer: Resolved workspace_id: {workspace_id_to_use}")
else:
workspace_id_to_use = data.workspace_id
if not data.workspace_slug:
if workspace_id_to_use:
resolved_workspace_slug = await resolve_workspace_slug(workspace_id_to_use, data.workspace_slug)
else:
log.warning("Mobile Queue-answer: No workspace_id to resolve workspace_slug from")
resolved_workspace_slug = None
else:
resolved_workspace_slug = data.workspace_slug
from uuid import uuid4
token_id = uuid4()
user_message_res = await upsert_message(
message_id=token_id,
chat_id=data.chat_id, # type: ignore[arg-type]
content=data.query,
parsed_content=None,
user_type=UserTypeChoices.USER.value,
llm_model=data.llm,
workspace_slug=resolved_workspace_slug,
db=db,
)
if user_message_res.get("message") != "success":
return JSONResponse(status_code=500, content={"detail": "Failed to create message"})
# 2. Stash the full ChatRequest inside a flow-step row
from fastapi.encoders import jsonable_encoder
flow_step_payload = {
"step_order": 0,
"step_type": FlowStepType.TOOL.value,
"tool_name": "QUEUE",
"content": "queued chat request",
# Use FastAPI's encoder to turn UUIDs/datetimes into JSON-serialisable primitives
"execution_data": jsonable_encoder(data),
"is_planned": False, # QUEUE is not a planned action, it's an internal tool
"is_executed": False, # QUEUE is not executed by user
}
flow_res = await upsert_message_flow_steps(message_id=token_id, chat_id=data.chat_id, db=db, flow_steps=[flow_step_payload])
if flow_res.get("message") != "success":
return JSONResponse(status_code=500, content={"detail": "Failed to queue request"})
# Backfill chat record with workspace details (after successful message creation)
try:
from sqlalchemy import select
from pi.app.models.chat import Chat
# Get the current chat record to preserve existing data
chat_stmt = select(Chat).where(Chat.id == data.chat_id) # type: ignore[arg-type]
chat_result = await db.execute(chat_stmt)
existing_chat = chat_result.scalar_one_or_none()
if existing_chat:
# Update the chat with workspace details
if workspace_id_to_use is not None:
existing_chat.workspace_id = workspace_id_to_use
if resolved_workspace_slug is not None:
existing_chat.workspace_slug = resolved_workspace_slug
# Always update workspace_in_context as it's a required field in ChatRequest
existing_chat.workspace_in_context = data.workspace_in_context
await db.commit()
else:
log.warning(f"Chat {data.chat_id} not found for backfill")
except Exception as e:
log.error(f"Error backfilling chat workspace details: {e}")
# Don't fail the request if backfill fails
return {"stream_token": str(token_id)}
@mobile_router.get("/stream-answer/{token}")
async def stream_answer(
token: UUID4,
db: AsyncSession = Depends(get_async_session),
auth_token: HTTPAuthorizationCredentials = Depends(jwt_schema),
):
"""Second phase of two-step flow.
Looks up the queued ChatRequest by token (message_id), deletes the queue
entry, and then re-uses the existing /get-answer/ logic to start the SSE
stream via a pure GET endpoint.
Also handles REGENERATE when called with an existing user message ID
(no QUEUE flow step present)."""
try:
auth = await validate_jwt_token(auth_token)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
user_id = auth.user.id
except Exception:
return JSONResponse(status_code=401, content={"detail": "Invalid JWT"})
# Locate the queued flow step
from sqlalchemy import select
stmt = (
select(MessageFlowStep)
.where(MessageFlowStep.message_id == token) # type: ignore[arg-type]
.where(MessageFlowStep.tool_name == "QUEUE") # type: ignore[arg-type]
)
res = await db.execute(stmt)
flow_step: MessageFlowStep | None = res.scalar_one_or_none()
if not flow_step:
# REGENERATE FLOW: No QUEUE flow step means this is regenerate
log.info(f"Mobile: No QUEUE flow step found for token {token}. Treating as regenerate request.")
# 1. Get the user message
user_message = await get_message_by_id(db, token)
if not user_message:
log.warning(f"Mobile: Message not found for token {token}")
return JSONResponse(status_code=404, content={"detail": "Message not found"})
# 2. Mark old assistant response as replaced BEFORE generating new one
# This ensures retrieve_chat_history() won't include it in context for new generation
marked = await mark_assistant_response_as_replaced(db, user_message.id)
if marked:
log.info(f"Mobile: Marked old assistant response(s) as replaced for message {token}")
else:
log.info(f"Mobile: No existing assistant response found for message {token} (first generation or already replaced)")
# 4. Reconstruct ChatRequest from user message
try:
queued_request = await reconstruct_chat_request_from_message(db, user_message, user_id)
except Exception as e:
log.error(f"Mobile: Error reconstructing ChatRequest from message {token}: {e}")
return JSONResponse(status_code=500, content={"detail": "Failed to reconstruct request"})
# 5. Pass token_id so new assistant message reuses same user message as parent
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Mobile: Failed to attach token_id to regenerate request context: {e!s}")
log.info(f"Mobile: Regenerating response for message {token}")
# 6. Stream new response (get_answer will call process_query_stream)
# When process_query_stream calls retrieve_chat_history,
# it will NOT see the old response because is_replaced=True
return await get_answer(data=queued_request, token=auth_token, db=db)
else:
# NORMAL FLOW: QUEUE flow step exists, this is first generation
# Parse the stored ChatRequest
try:
raw_data = flow_step.execution_data or {}
# Check if this is an OAuth message (has oauth_required field)
if flow_step.oauth_required:
# This is an OAuth message - check if OAuth is now complete
# by looking at the oauth_completed column
if flow_step.oauth_completed:
# OAuth is now complete! Process the request normally
log.info(f"Mobile: OAuth completed for message {token}. Processing request.")
# Control continues to the normal processing code below
else:
# OAuth is still required
return JSONResponse(
status_code=401,
content={
"detail": "OAuth authorization still required. Please complete OAuth authentication first.",
"error_code": "OAUTH_REQUIRED",
},
)
# Convert empty-string UUIDs to None so pydantic validation passes
for field in ["project_id", "workspace_id", "chat_id"]:
if field in raw_data and raw_data[field] == "":
raw_data[field] = None
raw_data["user_id"] = user_id
queued_request = ChatRequestMobile.parse_obj(raw_data)
except Exception as e:
log.error(f"Mobile: Malformed execution_data for token {token}: {e!s}")
return JSONResponse(status_code=500, content={"detail": "Corrupted queued request"})
# Consume the queue entry so the token is single-use (for normal flow)
await db.delete(flow_step)
await db.commit()
# Pass the token/message_id forward so downstream processing reuses the same Message row
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Mobile: Failed to attach token_id to queued request context: {e!s}")
# Delegate to existing get_answer for streaming
return await get_answer(data=queued_request, token=auth_token, db=db)
@mobile_router.delete("/delete-chat/")
async def delete_chat(
data: DeleteChatRequestMobile,
apps/pi/pi/app/models/message.py
+8
View File
@@ -58,6 +58,14 @@ class Message(BaseModel, table=True):
default=None,
)
# Regenerate/versioning fields
is_replaced: bool = Field(
default=False,
nullable=False,
description="Whether this message was replaced by regeneration (hidden from history)",
sa_column_kwargs={"server_default": text("false")},
)
# Relationships
message_feedbacks: List["MessageFeedback"] = Relationship(back_populates="message", sa_relationship_kwargs={"lazy": "selectin"})
apps/pi/pi/app/schemas/mobile/chat.py
+2
View File
@@ -1,4 +1,5 @@
from pi.app.schemas.chat import ChatFeedback
from pi.app.schemas.chat import ChatInitializationRequest
from pi.app.schemas.chat import ChatRequest
from pi.app.schemas.chat import ChatSearchRequest
from pi.app.schemas.chat import ChatSearchResponse
@@ -20,6 +21,7 @@ from pi.app.schemas.chat import TitleRequest
from pi.app.schemas.chat import UnfavoriteChatRequest
ChatRequestMobile = ChatRequest
ChatInitializationRequestMobile = ChatInitializationRequest
TitleRequestMobile = TitleRequest
GetThreadsMobile = GetThreads
ChatTypeMobile = ChatType
apps/pi/pi/manage.py
+3 -3
View File
@@ -239,7 +239,7 @@ def run_celery_worker(
"--without-gossip",
]
typer.echo(f"Starting Celery worker with command: {" ".join(cmd)}")
typer.echo(f'Starting Celery worker with command: {" ".join(cmd)}')
subprocess.run(cmd)
@@ -265,7 +265,7 @@ def run_celery_beat(
schedule_file,
]
typer.echo(f"Starting Celery Beat with command: {" ".join(cmd)}")
typer.echo(f'Starting Celery Beat with command: {" ".join(cmd)}')
subprocess.run(cmd)
@@ -291,7 +291,7 @@ def run_celery_flower(
address,
]
typer.echo(f"Starting Celery Flower with command: {" ".join(cmd)}")
typer.echo(f'Starting Celery Flower with command: {" ".join(cmd)}')
typer.echo(f"Flower will be available at http://{address}:{port}")
subprocess.run(cmd)
apps/pi/pi/scripts/celery_runner.py
+2 -2
View File
@@ -39,7 +39,7 @@ def run_celery_worker(
"--without-gossip",
]
typer.echo(f"Starting Celery worker with command: {" ".join(cmd)}")
typer.echo(f'Starting Celery worker with command: {" ".join(cmd)}')
subprocess.run(cmd)
@@ -62,7 +62,7 @@ def run_celery_beat(
schedule_file,
]
typer.echo(f"Starting Celery Beat with command: {" ".join(cmd)}")
typer.echo(f'Starting Celery Beat with command: {" ".join(cmd)}')
subprocess.run(cmd)
apps/pi/pi/services/retrievers/pg_store/chat.py
+9 -3
View File
@@ -387,8 +387,13 @@ async def retrieve_chat_history(
user_chat_preference_result = await db.execute(user_chat_preference_query)
user_chat_preference = user_chat_preference_result.scalar_one_or_none()
# Step 4: Get messages for this chat ordered by sequence
message_query = select(Message).where(Message.chat_id == chat_id).order_by(Message.sequence) # type: ignore[arg-type]
# Step 4: Get messages for this chat ordered by sequence (exclude replaced messages)
message_query = (
select(Message)
.where(Message.chat_id == chat_id) # type: ignore[arg-type]
.where(~Message.is_replaced) # type: ignore[arg-type]
.order_by(Message.sequence) # type: ignore[arg-type]
)
message_result = await db.execute(message_query)
messages = list(message_result.scalars().all())
@@ -398,11 +403,12 @@ async def retrieve_chat_history(
message_flow_steps_result = await db.execute(message_flow_steps_query)
message_flow_steps = list(message_flow_steps_result.scalars().all())
# Step 5: Get most recent assistant message to extract LLM model
# Step 5: Get most recent assistant message to extract LLM model (exclude replaced messages)
last_assistant_message_query = (
select(Message)
.where(Message.chat_id == chat_id) # type: ignore[arg-type]
.where(Message.user_type == UserTypeChoices.ASSISTANT.value) # type: ignore[arg-type]
.where(~Message.is_replaced) # type: ignore[arg-type]
.order_by(desc(Message.created_at)) # type: ignore[arg-type]
.limit(1)
)
apps/pi/pi/services/retrievers/pg_store/message.py
+135 -3
View File
@@ -13,6 +13,7 @@ from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.models import Chat
from pi.app.models import LlmModelPricing
from pi.app.models import Message
from pi.app.models import MessageFeedback
@@ -23,6 +24,9 @@ from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageFeedbackTypeChoices
from pi.app.models.enums import MessageMetaStepType
from pi.app.models.enums import UserTypeChoices
from pi.app.models.message_attachment import MessageAttachment
from pi.app.schemas.chat import ChatRequest
from pi.config import LLMModels
from pi.services.retrievers.pg_store.json_serializer import sanitize_execution_data
log = logger.getChild(__name__)
@@ -163,12 +167,17 @@ async def update_message_feedback(
async def get_chat_messages(chat_id: UUID4, db: AsyncSession) -> Union[List[Message], Tuple[int, Dict[str, str]]]:
"""
Retrieves all messages for a chat ordered by sequence.
Retrieves all messages for a chat ordered by sequence (excluding replaced messages).
Returns either a list of messages (success) or a tuple of (status_code, response_content) for errors
"""
try:
# Get messages for this chat ordered by sequence
messages_query = select(Message).where(Message.chat_id == chat_id).order_by(Message.sequence) # type: ignore[arg-type]
# Get messages for this chat ordered by sequence (exclude replaced messages)
messages_query = (
select(Message)
.where(Message.chat_id == chat_id) # type: ignore[arg-type]
.where(~Message.is_replaced) # type: ignore[arg-type]
.order_by(Message.sequence) # type: ignore[arg-type]
)
result = await db.execute(messages_query)
messages = list(result.scalars().all()) # Convert to list for correct typing
@@ -521,3 +530,126 @@ async def upsert_message_meta(
await db.rollback()
log.error(f"create_message_meta failed. message_id={message_id}, error={e}")
return {"message": "error", "error": str(e)}
async def mark_assistant_response_as_replaced(db: AsyncSession, user_message_id: UUID4) -> bool:
"""
Mark the assistant response(s) for a user message as replaced.
Args:
db: Database session
user_message_id: The user message ID whose assistant responses should be marked
Returns:
True if any messages were marked, False otherwise
"""
try:
# First, get the user message to find its sequence and chat_id
user_msg_stmt = select(Message).where(Message.id == user_message_id) # type: ignore[arg-type]
user_msg_result = await db.execute(user_msg_stmt)
user_message = user_msg_result.scalar_one_or_none()
if not user_message:
log.warning(f"User message {user_message_id} not found")
return False
# Find assistant message(s) that come after this user message (by sequence)
# in the same chat and are not already replaced
stmt = (
select(Message)
.where(Message.chat_id == user_message.chat_id) # type: ignore[arg-type]
.where(Message.sequence > user_message.sequence) # type: ignore[arg-type]
.where(Message.user_type == UserTypeChoices.ASSISTANT.value) # type: ignore[arg-type]
.where(~Message.is_replaced) # type: ignore[arg-type]
.order_by(Message.sequence) # type: ignore[arg-type]
.limit(1) # Only mark the immediate next assistant response
)
result = await db.execute(stmt)
old_assistant_message = result.scalar_one_or_none()
if not old_assistant_message:
log.info(f"No assistant message found to mark as replaced for user message {user_message_id}")
return False
# Mark as replaced
old_assistant_message.is_replaced = True
await db.commit()
log.info(
f"Marked assistant message {old_assistant_message.id} (sequence {old_assistant_message.sequence}) "
f"as replaced for user message {user_message_id}"
)
return True
except Exception as e:
log.error(f"Error marking assistant response as replaced: {e}")
await db.rollback()
return False
async def get_message_by_id(db: AsyncSession, message_id: UUID4) -> Optional[Message]:
"""
Get a message by ID.
Args:
db: Database session
message_id: Message ID to retrieve
Returns:
Message object if found, None otherwise
"""
try:
stmt = select(Message).where(Message.id == message_id) # type: ignore[arg-type]
result = await db.execute(stmt)
return result.scalar_one_or_none()
except Exception as e:
log.error(f"Error getting message {message_id}: {e}")
return None
async def reconstruct_chat_request_from_message(db: AsyncSession, user_message: Message, user_id: UUID4):
"""
Reconstruct a ChatRequest from an existing user message for regeneration.
Args:
db: Database session
user_message: The user message to reconstruct from
user_id: User ID making the request
Returns:
ChatRequest object ready for streaming
"""
# Get chat details
chat_stmt = select(Chat).where(Chat.id == user_message.chat_id) # type: ignore[arg-type]
chat_result = await db.execute(chat_stmt)
chat = chat_result.scalar_one_or_none()
# Get attachments for this message
attachments_stmt = (
select(MessageAttachment)
.where(MessageAttachment.message_id == user_message.id) # type: ignore[arg-type]
.where(MessageAttachment.status == "uploaded") # type: ignore[arg-type]
.where(MessageAttachment.deleted_at.is_(None)) # type: ignore[union-attr]
)
attachments_result = await db.execute(attachments_stmt)
attachments = attachments_result.scalars().all()
attachment_ids = [att.attachment_id for att in attachments] if attachments else []
# Reconstruct ChatRequest
return ChatRequest(
query=user_message.content or "",
chat_id=user_message.chat_id,
user_id=user_id,
llm=user_message.llm_model or LLMModels.DEFAULT,
is_new=False, # Always false for regenerate
is_temp=False,
workspace_id=chat.workspace_id if chat else None,
workspace_slug=chat.workspace_slug if chat else None,
workspace_in_context=chat.workspace_in_context if chat else False,
is_project_chat=chat.is_project_chat if chat else False,
project_id=None, # Will be resolved from context if needed
attachment_ids=attachment_ids,
context={"token_id": str(user_message.id)}, # Reuse same message ID
source="web",
)
apps/web/ee/components/pi-chat/conversation/ai-message.tsx
+2
View File
@@ -94,6 +94,8 @@ export const AiMessage = observer((props: TProps) => {
id={id}
workspaceId={workspaceId}
feedback={feedback}
queryId={query_id}
isLatest={!!isLatest}
handleConvertToPage={handleConvertToPage}
/>
)}
apps/web/ee/components/pi-chat/conversation/feedback.tsx
+21 -7
View File
@@ -1,5 +1,5 @@
import { useState } from "react";
import { Copy, FilePlus2, ThumbsDown, ThumbsUp } from "lucide-react";
import { Copy, FilePlus2, ThumbsDown, ThumbsUp, Repeat2 } from "lucide-react";
import { setToast, TOAST_TYPE } from "@plane/propel/toast";
import { cn, Tooltip } from "@plane/ui";
import { copyTextToClipboard } from "@plane/utils";
@@ -13,16 +13,18 @@ export type TProps = {
id: string;
workspaceId: string | undefined;
feedback: EFeedback | undefined;
queryId: string | undefined;
isLatest: boolean;
handleConvertToPage?: () => void;
};
export const Feedback = (props: TProps) => {
// props
const { answer, activeChatId, id, workspaceId, feedback, handleConvertToPage } = props;
const { answer, activeChatId, id, workspaceId, feedback, queryId, isLatest, handleConvertToPage } = props;
// states
const [isFeedbackModalOpen, setIsFeedbackModalOpen] = useState(false);
// store
const { sendFeedback } = usePiChat();
const { sendFeedback, regenerateAnswer } = usePiChat();
// handlers
const handleCopyLink = () => {
copyTextToClipboard(answer).then(() => {
@@ -49,6 +51,14 @@ export const Feedback = (props: TProps) => {
});
}
};
const handleRewrite = async () => {
try {
if (!queryId || !workspaceId) return;
await regenerateAnswer(activeChatId, queryId, workspaceId);
} catch (e) {
console.error(e);
}
};
return (
<div className="flex gap-4 mt-6">
@@ -98,10 +108,14 @@ export const Feedback = (props: TProps) => {
onSubmit={(feedbackMessage) => handleFeedback(EFeedback.NEGATIVE, feedbackMessage)}
/>
{/* Rewrite will be available in the future */}
{/* <div className="flex text-sm font-medium gap-1 cursor-pointer">
<Repeat2 size={20} onClick={() => console.log()} className="my-auto cursor-pointer" /> Rewrite
</div> */}
{/* Rewrite */}
{isLatest && (
<Tooltip tooltipContent="Rewrite" position="bottom" className="mb-4">
<button onClick={handleRewrite}>
<Repeat2 strokeWidth={1.5} size={20} className="my-auto text-pi-700 transition-colors" />
</button>
</Tooltip>
)}
{/* Convert to page */}
<div className="flex text-sm font-medium gap-1 cursor-pointer">