Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion backend/backend/application/context/chat_ai_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -326,9 +326,18 @@ def _process_completed(self, *args, **kwargs):
# On successful completion, Closing the event thread
raise StopIteration

def _process_chat_intent(self, *args, **kwargs):
send_socket_message(
sid=kwargs["sid"],
channel_id=kwargs["channel_id"],
chat_id=kwargs["chat_id"],
chat_message_id=kwargs["chat_message_id"],
chat_intent_name=kwargs.get("chat_intent"),
)

def process_event(self, *args, **kwargs):
supported_events = ["thought_chain", "prompt_response",
"summary", "chat_name", "completed"]
"summary", "chat_name", "completed", "chat_intent"]
event_type = kwargs.get("event_type")
if event_type not in supported_events:
raise ValueError(f"Unsupported event type: {event_type}")
Expand Down
12 changes: 2 additions & 10 deletions backend/backend/application/context/chat_message_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,37 +144,30 @@ def persist_prompt(
llm_model_architect: str,
llm_model_developer: str,
generated_chat_res_id: str = None,
chat_intent_id: str = None,
chat_id: str = None,
user=None,
) -> ChatMessage:
"""
Create a new prompt within a Chat. If chat_id is None, create a new Chat.
Return the chat_message_id of the newly created ChatMessage.
chat_intent is left null here; the AI service auto-detects it and the
backend persists it on the message when the response arrives.
"""
if not prompt.strip():
raise InvalidChatPrompt()

chat_intent = None
transformation_type = 'TRANSFORM' if discussion_type == 'GENERATE' else 'DISCUSSION'
if chat_intent_id:
try:
chat_intent = ChatIntent.objects.get(chat_intent_id=chat_intent_id)
except ChatIntent.DoesNotExist:
chat_intent = None

if not chat_id:
chat = Chat.objects.create(
project=self.project_instance,
chat_name="Untitled Chat",
chat_intent=chat_intent,
llm_model_architect=llm_model_architect,
llm_model_developer=llm_model_developer,
user=user,
)
else:
chat = self._get_chat_or_raise(chat_id=chat_id, must_be_active=True)
chat.chat_intent = chat_intent
chat.llm_model_architect = llm_model_architect
chat.llm_model_developer = llm_model_developer
chat.discussion_type = discussion_type
Expand All @@ -185,7 +178,6 @@ def persist_prompt(
chat_message = ChatMessage.objects.create(
chat=chat,
prompt=prompt,
chat_intent=chat_intent,
llm_model_architect=llm_model_architect,
llm_model_developer=llm_model_developer,
discussion_type= discussion_type,
Expand Down
64 changes: 48 additions & 16 deletions backend/backend/application/context/llm_context.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import json
import logging
from typing import Any
from typing import Any, Optional

import eventlet
import redis
Expand Down Expand Up @@ -52,12 +52,47 @@ def create_redis_xgroup(self, channel_id, group_id):
else:
raise

def _resolve_chat_intent(self, chat_message_id: str, data: dict, content: Any) -> Optional[str]:
"""
Pull the AI-detected chat_intent out of the inbound payload and persist it
to ChatMessage.chat_intent the first time we see it for a given message.
Subsequent events re-use the cached value.
"""
if not hasattr(self, "_chat_intent_by_msg"):
self._chat_intent_by_msg = {}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cleaner to initialize self._chat_intent_by_msg = {} in __init__ rather than lazy-init via hasattr. If LLMServerContext is ever pooled or reused across requests (the rest of the class is stateful enough that this seems plausible), this dict carries entries across requests because there's no reset point. An explicit init in __init__ makes the lifecycle obvious.


cached = self._chat_intent_by_msg.get(chat_message_id)
if cached:
return cached

intent_name = data.get("chat_intent")
if not intent_name and isinstance(content, dict):
intent_name = content.get("chat_intent")
if not intent_name:
return None

try:
from backend.core.models.chat_intent import ChatIntent
from backend.core.models.chat_message import ChatMessage
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These imports run on every stream event (the helper is invoked from process_message for all event types, not just chat_intent). Python's import cache makes the cost small but non-zero, and there's no circular-import excuse here — both modules are already top-level imported elsewhere. Hoist to the top of the file.

chat_intent = ChatIntent.objects.get(name=intent_name)
chat_message = ChatMessage.objects.get(chat_message_id=chat_message_id)
chat_message.chat_intent = chat_intent
chat_message.save(update_fields=["chat_intent"])
# Mirror the pre-PR behavior where Chat.chat_intent tracked the
# latest message's intent (used by PastConversations badge).
chat_message.chat.chat_intent = chat_intent
chat_message.chat.save(update_fields=["chat_intent"])
except Exception as e:
logging.error(f"Failed to persist chat_intent={intent_name}: {e}")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bare except Exception masks distinct failure modes.

This catches ChatIntent.DoesNotExist (typo / unseeded value), OperationalError (transient DB), IntegrityError, etc. all the same. The frontend then renders a response with no Apply / Run SQL footer (because Conversation.jsx returns null when chat_intent_name is missing), and the only signal is one logging.error line. For something that controls billing routing and the primary CTA on every AI response, that's too quiet — narrow the except, and at minimum surface a sentry/metric so we notice when the AI starts emitting an unknown intent name.

See companion comment on the line below for the cache-write-outside-try issue, which compounds this.


self._chat_intent_by_msg[chat_message_id] = intent_name
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cache write outside the try — silently locks in failure.

If the ChatIntent.objects.get / chat_message.save block above raised, the bare except Exception swallowed it and we still reach this line and record intent_name in the cache. Subsequent events on the same stream then short-circuit on cached, never retrying the persist. Meanwhile the socket push at chat_ai_context._process_chat_intent does still send the name to the frontend (it reads from kwargs["chat_intent"] which is the cached return value), so the UI shows the intent during streaming but the DB row stays null forever — historical reload is then inconsistent with what the user just saw.

Move the cache write into the try block (after the save), and narrow the except so ChatIntent.DoesNotExist (typo / unseeded value) doesn't get conflated with transient OperationalErrors.

Suggested shape:

try:
    chat_intent_obj = ChatIntent.objects.get(name=intent_name)
except ChatIntent.DoesNotExist:
    logging.error(f"Unknown chat_intent name from AI: {intent_name!r}")
    return None  # don't cache; let a later event try again

try:
    with transaction.atomic():
        ChatMessage.objects.filter(chat_message_id=chat_message_id) \
            .update(chat_intent=chat_intent_obj)
        Chat.objects.filter(chat_messages__chat_message_id=chat_message_id) \
            .update(chat_intent=chat_intent_obj)
except DatabaseError:
    logging.exception("Failed to persist chat_intent")
    return None

self._chat_intent_by_msg[chat_message_id] = intent_name
return intent_name

return intent_name

def process_message(
self,
sid: str,
channel_id: str,
chat_id: str,
chat_intent: str,
payload: dict[str, Any],
discussion_status: str
):
Expand All @@ -75,6 +110,7 @@ def process_message(
2: "summary",
3: "chat_name",
4: "completed",
5: "chat_intent",
99: "stop",
}

Expand All @@ -83,6 +119,8 @@ def process_message(
chat_message_id = data["chat_message_id"]
content = data["content"]

chat_intent = self._resolve_chat_intent(chat_message_id, data, content)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Retry no longer signals "this is a retry" to the AI gatekeeper.

Pre-PR, the is_retry branch in process_prompt set chat_intent = ChatMessageStatus.TRANSFORM_RETRY and forwarded it in llm_payload. Post-PR the assignment is gone and chat_intent is no longer in the outbound payload at all (see process_prompt below). The AI gatekeeper now classifies the retry prompt — which is literally "Faulty yaml: <yaml> \n Error: <error>" — from scratch. There's a real chance it routes that to SQL or INFO instead of TRANSFORM, breaking retry.

Second, related concern: when the gatekeeper does classify the retry, the new CHAT_INTENT event will hit this _resolve_chat_intent and overwrite ChatMessage.chat_intent. If the new classification differs from the original message's intent, the row's intent flips on retry — surprising for anyone reading historical data, and breaks the ResponseFooter rendering for the original message.

Mitigations (pick one):

  • Keep an explicit is_retry: True field in llm_payload and have the AI side short-circuit gatekeeper when set
  • Or have _resolve_chat_intent not overwrite ChatMessage.chat_intent once it's already set
  • And add retry coverage to the test plan — there's none right now


if event_type == "chat_name":
self.chat_name = data["content"]
self.persist_response(
Expand Down Expand Up @@ -120,7 +158,7 @@ def _validate_message(self, group_id, channel_id):
)
return messages

def _handle_redis_message(self, sid, channel_id, chat_id, chat_intent, group_id, messages, discussion_status: str):
def _handle_redis_message(self, sid, channel_id, chat_id, group_id, messages, discussion_status: str):
for _, msg_list in messages:
for message_id, payload in msg_list:
logging.info(f" === Message ID: {message_id} ===")
Expand All @@ -129,7 +167,6 @@ def _handle_redis_message(self, sid, channel_id, chat_id, chat_intent, group_id,
sid=sid,
channel_id=channel_id,
chat_id=chat_id,
chat_intent=chat_intent,
payload=payload,
discussion_status=discussion_status
)
Expand All @@ -138,7 +175,7 @@ def _handle_redis_message(self, sid, channel_id, chat_id, chat_intent, group_id,
self.redis_client.xack(channel_id, group_id, message_id)

def __stream_listener(
self, sid: str, channel_id: str, chat_id: str, chat_message_id: str, chat_intent: str, group_id: str, discussion_status: str
self, sid: str, channel_id: str, chat_id: str, chat_message_id: str, group_id: str, discussion_status: str
):

while True:
Expand All @@ -148,7 +185,7 @@ def __stream_listener(
if not messages:
continue

self._handle_redis_message(sid, channel_id, chat_id, chat_intent, group_id, messages, discussion_status)
self._handle_redis_message(sid, channel_id, chat_id, group_id, messages, discussion_status)

except redis.exceptions.RedisError as e:
logging.error(f"[REDIS ERROR] {e}")
Expand Down Expand Up @@ -196,15 +233,15 @@ def __stream_listener(
)
break

def listen_to_redis_stream(self, sid: str, channel_id: str, chat_id: str, chat_message_id: str, chat_intent: str, discussion_status: str):
def listen_to_redis_stream(self, sid: str, channel_id: str, chat_id: str, chat_message_id: str, discussion_status: str):
"""Listens to the Redis stream from llm server and processes the messages."""
group_id = f"group_{chat_id}_{chat_message_id}"
self.create_redis_xgroup(channel_id, group_id)
self.__stream_listener(sid, channel_id, chat_id, chat_message_id, chat_intent, group_id, discussion_status)
self.__stream_listener(sid, channel_id, chat_id, chat_message_id, group_id, discussion_status)

def stream_prompt_response(self, sid: str, channel_id: str, chat_id: str, chat_message_id: str, chat_intent: str, discussion_status: str):
def stream_prompt_response(self, sid: str, channel_id: str, chat_id: str, chat_message_id: str, discussion_status: str):
"""Starts a background thread to listen redis pubsub channel from AI server"""
args = (sid, channel_id, chat_id, chat_message_id, chat_intent, discussion_status)
args = (sid, channel_id, chat_id, chat_message_id, discussion_status)
try:
sio.start_background_task(self.listen_to_redis_stream, *args)
except Exception as e:
Expand Down Expand Up @@ -237,12 +274,10 @@ def process_prompt(self, sid: str, channel_id: str, chat_id: str, chat_message_i
"GENERATE": ChatMessageStatus.GENERATE,
}
if is_retry:
chat_intent = ChatMessageStatus.TRANSFORM_RETRY
prompt = (
f"Faulty yaml:{chat_message.technical_content} \n Error:{chat_message.transformation_error_message}"
)
else:
chat_intent = chat_message.chat_intent.name
prompt = chat_message.prompt

if discussion_status in DISCUSSION_STATUS_MAP:
Expand Down Expand Up @@ -311,7 +346,6 @@ def process_prompt(self, sid: str, channel_id: str, chat_id: str, chat_message_i
"db_map": db_metadata,
"visitran_model": visitran_models,
"chat_name": chat_name,
"chat_intent": chat_intent,
"db_type": self.project_instance.database_type,
"llm_model_architect": chat_message.llm_model_architect,
"llm_model_developer": chat_message.llm_model_developer,
Expand All @@ -335,7 +369,6 @@ def process_prompt(self, sid: str, channel_id: str, chat_id: str, chat_message_i
channel_id=channel_id,
chat_id=chat_id,
chat_message_id=chat_message_id,
chat_intent=chat_intent,
discussion_status=chat_message.discussion_type,
)

Expand All @@ -347,10 +380,9 @@ def process_prompt(self, sid: str, channel_id: str, chat_id: str, chat_message_i
channel_id=channel_id,
chat_id=chat_id,
chat_message_id=chat_message_id,
chat_intent=chat_intent,
discussion_status=chat_message.discussion_type,
)
logging.info(f"process_prompt: chat_intent={chat_intent}, sid={sid}, channel_id={channel_id}")
logging.info(f"process_prompt: sid={sid}, channel_id={channel_id}")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Likely billing bug — TRANSFORM may be double-charged in non-WS mode.

In the cloud override of process_prompt (pluggable_apps/subscriptions/context.py), right after super().process_prompt(...) returns we do:

chat_intent = chat_message.chat_intent.name if chat_message.chat_intent else "INFO"
if chat_intent != "TRANSFORM":
    self.consume_tokens(...)

In this (non-WS) branch the Redis stream listener is started after send_event_to_llm_server(...) returns, so when super().process_prompt returns the listener has just begun and chat_message.chat_intent is still None. That falls through to "INFO" != "TRANSFORM"consume_tokens() runs at generation. Then transformation_save charges again on Apply success. Net: every TRANSFORM in non-WS mode is charged twice.

WS mode probably happens to work because eventlet yields on Redis I/O and the listener drains in time, but relying on yield ordering for billing correctness is brittle.

Fix options, in order of preference:

  1. Have the AI side guarantee CHAT_INTENT (event_type=5) is the first event, and drain it synchronously inside process_prompt before returning → no race in either mode.
  2. Move the cloud-side consume_tokens decision to a hook fired from _resolve_chat_intent (i.e. when intent is known), not from process_prompt's return path.
  3. At minimum, make the cloud override re-fetch and skip charging when chat_intent is still null (defer to Apply or to a later event), rather than defaulting to "INFO".

The test plan claims billing was verified — please add an explicit non-WS test before merge.

chat_message = self._get_chat_message(chat_id=chat_id, chat_message_id=chat_message_id)

return chat_message
Expand Down
4 changes: 4 additions & 0 deletions backend/backend/core/routers/chat/serializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@

class ChatSerializer(serializers.ModelSerializer):
user = UserMinimalSerializer(read_only=True)
chat_intent_name = serializers.CharField(
source='chat_intent.display_name', read_only=True, default=None
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same field name (chat_intent_name), different source from the chat-message serializer:

  • here: source='chat_intent.display_name' → e.g. "Transform"
  • in chat_message_serializer.py:14: source='chat_intent.name' → e.g. "TRANSFORM"

The rationale (chat list shows the human label, message-level needs the canonical token for branching) is fine, but the asymmetry under one shared field name is undocumented and a future bug magnet — anyone wiring up a new component will assume both return the same thing. Either rename the chat-list one to chat_intent_label (or similar) or pick a single source and have the frontend map for display.

)

class Meta:
model = Chat
Expand All @@ -13,6 +16,7 @@ class Meta:
'project_id',
'chat_name',
'chat_intent',
'chat_intent_name',
'created_at',
'modified_at',
'is_deleted',
Expand Down
19 changes: 5 additions & 14 deletions backend/backend/core/routers/chat/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,6 @@ def persist_prompt(self, request: Request, project_id: str, *args, **kwargs) ->
data = request.data
chat_id = data.get("chat_id")
prompt = data.get("prompt")
chat_intent_id = data.get("chat_intent_id")
llm_model_architect = data.get("llm_model_architect")
llm_model_developer = data.get("llm_model_developer")
discussion_type = data.get('discussion_status')
Expand All @@ -139,26 +138,19 @@ def persist_prompt(self, request: Request, project_id: str, *args, **kwargs) ->
if discussion_type == "GENERATE":
generated_chat_res_id = data.get('final_discussion_id')

# Check token balance before processing the request
# Check token balance before processing the request.
# Intent is auto-detected by the AI service, so we don't know it yet —
# check against the worst-case (TRANSFORM) cost to avoid letting through
# a request the org can't afford.
try:
project = ProjectDetails.objects.get(project_uuid=project_id)
organization = project.organization

# Determine chat intent name for token calculation
chat_intent_name = "INFO" # Default
if chat_intent_id:
from backend.core.models.chat_intent import ChatIntent
try:
chat_intent = ChatIntent.objects.get(chat_intent_id=chat_intent_id)
chat_intent_name = chat_intent.name
except ChatIntent.DoesNotExist:
pass

self.fetch_token_balance(
llm_model_architect=llm_model_architect,
llm_model_developer=llm_model_developer,
organization=organization,
chat_intent_name=chat_intent_name
chat_intent_name="TRANSFORM"
)

except ProjectDetails.DoesNotExist:
Expand All @@ -172,7 +164,6 @@ def persist_prompt(self, request: Request, project_id: str, *args, **kwargs) ->
chat_message = chat_message_context.persist_prompt(
prompt=prompt,
chat_id=chat_id,
chat_intent_id=chat_intent_id,
llm_model_architect=llm_model_architect,
llm_model_developer=llm_model_developer,
discussion_type=discussion_type,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ class Meta:

class ChatMessageSerializer(serializers.ModelSerializer):
user = UserMinimalSerializer(read_only=True)
chat_intent_name = serializers.CharField(source='chat_intent.name', read_only=True, default=None)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See companion comment on chat/serializers.py:9. This one returns chat_intent.name ("TRANSFORM") while the chat-list serializer returns chat_intent.display_name ("Transform") under the same field name. Worth disambiguating.


class Meta:
model = ChatMessage
Expand All @@ -29,6 +30,7 @@ class Meta:
'transformation_status',
'transformation_error_message',
'chat_intent',
'chat_intent_name',
'llm_model_architect',
'llm_model_developer',
'created_at',
Expand Down
1 change: 1 addition & 0 deletions backend/backend/core/web_socket.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,7 @@ def send_socket_message(sid, channel_id, **kwargs):
"is_retry_transform",
"discussion_status",
"token_usage_data", # Add token usage data
"chat_intent_name",
]

unsupported_args = [arg for arg in kwargs.keys() if arg not in allowed_args]
Expand Down
8 changes: 3 additions & 5 deletions frontend/src/ide/chat-ai/ActionButtons.jsx
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@ const INFO_APPROVED =
const ActionButtons = memo(function ActionButtons({
chatMessageId,
savePrompt,
selectedChatIntent,
isLatestTransform,
uiAction,
message,
Expand Down Expand Up @@ -101,12 +100,12 @@ const ActionButtons = memo(function ActionButtons({
setTimeout(() => setIsOperationInProgress(false), 3000);

if (value === "GENERATE") {
savePrompt?.(label, selectedChatIntent, false, value, chatMessageId);
savePrompt?.(label, false, value, chatMessageId);
return;
}
savePrompt?.(label, selectedChatIntent, false, value);
savePrompt?.(label, false, value);
},
[savePrompt, selectedChatIntent, chatMessageId, isOperationInProgress]
[savePrompt, chatMessageId, isOperationInProgress]
);

const onApplyClick = useCallback(() => {
Expand Down Expand Up @@ -321,7 +320,6 @@ ActionButtons.displayName = "ActionButtons";
ActionButtons.propTypes = {
chatMessageId: PropTypes.string.isRequired,
savePrompt: PropTypes.func,
selectedChatIntent: PropTypes.string,
isLatestTransform: PropTypes.bool.isRequired,
uiAction: PropTypes.object,
message: PropTypes.object,
Expand Down
Loading
Loading