import json

from google.genai import types
from google.adk.runners import Runner
from google.adk.events import Event
from google.adk.events.event_actions import EventActions
from google.adk.sessions import Session
from google.adk.apps.app import App, EventsCompactionConfig
from google.adk.agents.context_cache_config import ContextCacheConfig
from langchain_core.output_parsers import PydanticOutputParser

from app.base.llm.common import BotLLM
from app.base.models import EnhancedAiSearchQueryModel
from app import app_settings as settings
from app.app_settings import logger
from app.base.choices import (
    TextGenerationModelChoices, AgentNameChoices,
    IntentionChoiceChoices
)
from app.base.helpers import TokenCounter
from app.base.choices import FeedSourceChoices
from app.base.recommendations.functions import (
    filter_products_by_recommendations, get_products_by_ids
)
from agent.helpers import agent_intention_mapper
from agent import format_instruction
from agent.root_agent.agent import create_root_agent
from agent.root_agent import instruction_prompt
from chat.schema import SessionExtendedModel
from chat.schema import (
    ChatResponseModel, ChatResponseContentModel,
    ChatResponseMetaDataModel
)


DEFAULT_ERROR_RESPONSE = "Hey, sorry something went wrong, please try again later."


class AgentClient:
    """
    Service class for interacting with the Assistant Agent.
    Usage:
        chat_service = await AgentClient.create(session_overall_info)
        await chat_service.chat()
    """

    def __init__(self, session_info: SessionExtendedModel):
        self.app_name = settings.ADK_APP_NAME
        self.session_info: SessionExtendedModel = session_info
        self.session_service = settings.ADK_DB_SERVICE
        self.user_id = f"user_{self.session_info.bot_id}"

        self.initial_state = {
            "org_id": session_info.org_id,
            "bot_id": session_info.bot_id,
            "raw_query": session_info.query,
            "customer_name": getattr(session_info, 'customer_name', ''),
            "customer_email": getattr(session_info, 'customer_email', ''),
            "customer_support_email": session_info.customer_support_email,
            "customer_support_number": session_info.customer_support_number,
            "language": session_info.language,
            "company_description": session_info.company_description,
            "require_human_assistance": session_info.require_human_assistance,
            "show_stock_info": session_info.show_stock_info,
            "support_online": session_info.is_admin_available,

            # POST processing values
            "order_email": None,
            "order_id": None,
            "user_reformed_query": None,
            "recent_messages": [],
            "recommended_product_ids": [],
            "recommended_variant_ids": []
        }

        # Set in _setup, used in chat()
        self.app = None
        self.runner = None

        self.token_counter = TokenCounter(TextGenerationModelChoices.GEMINI_3_FLASH_PREVIEW.value)

    @classmethod
    async def create(cls, session_info: SessionExtendedModel) -> "AgentClient":
        """Factory method — creates and fully initializes the AgentClient."""
        instance = cls(session_info)
        await instance._setup()
        return instance

    async def _setup(self):
        """
        Full initialization sequence:
          1. Get or create ADK session
          2. Build search query + write intention to session state
          3. Reload session to read intention
          4. Build root agent narrowed by intention
          5. Create App + Runner
        """

        # Step 1: get or create session
        await self._get_or_create_session(
            user_id=self.user_id,
            session_id=self.session_info.user_session_id
        )

        # Step 2: build search query — writes intention, search_query, etc. to session state
        await self.build_search_query()

        # Step 3: reload session so we can read the intention written in step 2
        active_session = await self.session_service.get_session(
            app_name=self.app_name,
            user_id=self.user_id,
            session_id=self.session_info.user_session_id
        )
        intention = active_session.state.get('intention', '').lower()
        logger.info(f"Resolved intention for agent routing: {intention}")

        # Step 4: build root agent with narrowed sub-agents based on intention
        agent = await create_root_agent(self.session_info, intention=intention)

        # Step 5: create App + Runner
        self.app = App(
            name=self.app_name,
            root_agent=agent,
            events_compaction_config=EventsCompactionConfig(
                compaction_interval=3,
                overlap_size=1
            ),
            context_cache_config=ContextCacheConfig(
                min_tokens=2048,    # Minimum tokens to trigger caching
                ttl_seconds=600,    # Store for up to 10 minutes
                cache_intervals=5,  # Refresh after 5 uses
            )
        )
        self.runner = Runner(app=self.app, session_service=self.session_service)

    async def _get_or_create_session(self, user_id: str, session_id: str) -> Session:
        """Get existing session or create a new one, syncing mutable fields if needed."""

        session: Session = await self.session_service.get_session(
            app_name=self.app_name,
            user_id=user_id,
            session_id=session_id
        )

        if session:
            # Fields that may change between requests
            updatable_fields = {
                "raw_query": self.session_info.query,
                "support_online": self.session_info.is_admin_available,
                "company_description": self.session_info.company_description,
                "require_human_assistance": self.session_info.require_human_assistance,
                "show_stock_info": self.session_info.show_stock_info,
                "customer_support_email": self.session_info.customer_support_email,
                "customer_support_number": self.session_info.customer_support_number,
            }

            # Only emit an event if something actually changed
            state_delta = {
                key: val
                for key, val in updatable_fields.items()
                if session.state.get(key) != val
            }

            if state_delta:
                logger.debug(f"Updating session state delta: {state_delta}")
                try:
                    await self.session_service.append_event(
                        session=session,
                        event=Event(author="user", actions=EventActions(state_delta=state_delta))
                    )
                except Exception as e:
                    pass
        else:
            session = await self.session_service.create_session(
                app_name=self.app_name,
                user_id=user_id,
                session_id=session_id,
                state=self.initial_state
            )

        return session

    async def build_search_query(self) -> None:
        """
        Call the LLM to build an enhanced search query from the raw user query.
        Writes search_query, intention, search_type, item_identifiers (and optionally
        yes_no_query) into session state.
        """

        session = await self.session_service.get_session(
            app_name=self.app_name,
            user_id=self.user_id,
            session_id=self.session_info.user_session_id
        )
        raw_query = session.state.get('raw_query', self.session_info.query)
        recent_messages = session.state.get('recent_messages', [])
        company_description = session.state.get('company_description', self.session_info.company_description)

        parser = PydanticOutputParser(pydantic_object=EnhancedAiSearchQueryModel)
        _instruction_prompt = instruction_prompt.conversation_query_build_prompt.format(
            query=raw_query,
            recent_messages=recent_messages,
            company_description=company_description,
        )

        _instruction_prompt += format_instruction.format(format_instruction=parser.get_format_instructions())

        llm = BotLLM(
            org_id=session.state.get('org_id', self.session_info.org_id),
            bot_id=session.state.get('bot_id', self.session_info.bot_id),
            query=raw_query,
            instruction_prompt=_instruction_prompt,
            parser=parser,
            thinking_level=types.ThinkingLevel.LOW,
        )
        response, token, cost, error = await llm.ask()

        search_query = response.get("search_query")
        intention = (response.get("intention") or "").lower()
        search_type = response.get("search_type")
        item_identifiers = response.get("item_identifiers", [])

        state_delta = {
            "search_query": search_query if (not error and search_query) else raw_query,
            "intention": intention or None,
            "search_type": search_type,
            "item_identifiers": item_identifiers,
        }

        # yes_no_query only relevant for product intention
        if intention == IntentionChoiceChoices.PRODUCT.value:
            state_delta["yes_no_query"] = response.get("yes_no_query")

        try:
            await self.session_service.append_event(
                session=session,
                event=Event(author="user", actions=EventActions(state_delta=state_delta))
            )
        except Exception as e:
            pass

    async def update_recent_messages(
        self,
        user_query: str, 
        ai_msg: str
    ) -> None:
        """Append the latest user/AI message pair to recent_messages in session state."""

        # Reload fresh session since runner has modified it
        session: Session = await self.session_service.get_session(
            app_name=self.app_name,
            user_id=self.user_id,
            session_id=self.session_info.user_session_id
        )

        recent_messages = session.state.get('recent_messages', [])
        recent_messages.append({"role": "user", "content": user_query})
        recent_messages.append({"role": "ai", "content": ai_msg})
        recent_messages = recent_messages[-30:]

        try:
            await self.session_service.append_event(
                session=session,
                event=Event(
                    author="user",
                    actions=EventActions(state_delta={"recent_messages": recent_messages})
                )
            )
        except Exception as e:
            pass

    async def _persist_recommended_product_ids(self, recommended_ids: list) -> None:
        """
        Merge newly recommended product/variant IDs into session state,
        keeping only the last 20 of each.
        """
        new_product_ids = [str(p.get("product_id")) for p in recommended_ids if p.get("product_id")]
        new_variant_ids = [str(v) for p in recommended_ids for v in (p.get("variant_ids") or [])]

        if not new_product_ids and not new_variant_ids:
            return

        try:
            session = await self.session_service.get_session(
                app_name=self.app_name,
                user_id=self.user_id,
                session_id=self.session_info.user_session_id
            )
            if not session:
                return

            existing_product_ids = session.state.get('recommended_product_ids', [])
            existing_variant_ids = session.state.get('recommended_variant_ids', [])

            updated_product_ids = [str(p) for p in existing_product_ids + new_product_ids if p][-20:]
            updated_variant_ids = [str(v) for v in existing_variant_ids + new_variant_ids if v][-20:]

            try:
                await self.session_service.append_event(
                    session=session,
                    event=Event(
                        author="user",
                        actions=EventActions(state_delta={
                        'recommended_product_ids': updated_product_ids,
                        'recommended_variant_ids': updated_variant_ids,
                        })
                    )
                )
            except Exception as e:
                pass

            logger.info(f"Persisted recommended product IDs: {updated_product_ids}")
            logger.info(f"Persisted recommended variant IDs: {updated_variant_ids}")

        except Exception as e:
            logger.error(f"Could not persist recommended product IDs: {e}", exc_info=True)

    async def chat(self) -> ChatResponseModel:
        """Process a chat message through the agent and return a structured response."""

        user_query = self.session_info.query
        user_id = self.user_id
        active_session = await self.session_service.get_session(
            app_name=self.app_name,
            user_id=user_id,
            session_id=self.session_info.user_session_id
        )

        # Tracking variables
        error = False
        agent_response = None
        tool_response_data = None
        tool_response_name = None
        last_event = None
        last_usage_metadata = None
        product_info = []
        query_intention = 'unknown'
        admin_notification = {'has_notification': False, 'notification_message': None}
        has_asked_for_order_status = False
        total_tokens = 0
        total_cost = 0

        try:
            new_message = types.Content(
                role="user",
                parts=[types.Part(text=user_query)]
            )

            # Run the agent
            async for event in self.runner.run_async(
                user_id=user_id,
                session_id=active_session.id,
                new_message=new_message
            ):
                last_event = event
                last_usage_metadata = event.usage_metadata
                event_content = event.content
                parts = event_content.parts if event_content else []

                for part in parts:
                    # CASE A: Final text response from an agent
                    if hasattr(part, 'text') and part.text:
                        try:
                            agent_response = json.loads(part.text)
                        except json.JSONDecodeError:
                            logger.error(f"Failed to parse JSON from agent response: {part.text}")
                            agent_response = {"ai_msg": part.text.strip()}

                    # CASE B: Tool (function) response — capture raw tool data
                    fr = getattr(part, 'function_response', None)
                    if fr and getattr(fr, 'response', None):
                        tool_response_data = fr.response
                        tool_response_name = fr.name

            # Token + cost accounting
            if last_usage_metadata:
                self.token_counter.add_from_metadata(last_usage_metadata)

            input_tokens = self.token_counter.get_input_tokens()
            output_tokens = self.token_counter.get_output_tokens()
            thoughts_tokens = self.token_counter.get_thoughts_tokens()
            total_tokens = self.token_counter.get_total_tokens()

            input_cost = self.token_counter.get_input_cost()
            thoughts_cost = self.token_counter.get_thoughts_cost()
            output_cost = self.token_counter.get_output_cost()
            total_cost = self.token_counter.get_total_cost()

            # Parse agent response
            author = last_event.author
            query_intention = agent_intention_mapper(author)
            ai_msg = agent_response.get('ai_msg', '')
            admin_notification = {
                "has_notification": agent_response.get("notify_human_assistant", False),
                "message": agent_response.get("notification_text", None)
            }
            has_asked_for_order_status = agent_response.get('has_asked_for_order_status', False)
            ai_response = {"ai_message": ai_msg, "product_info": product_info}

            logger.info(f"Last event author: {author}")
            logger.info(f"Query intention: {query_intention}")
            logger.info(f"Agent Response: {agent_response}")
            logger.info(f"Tool Name: {tool_response_name}")
            logger.info(f"Tool Response Data: {tool_response_data}")
            logger.info(
                f"\nToken Usage Costs:\n"
                f"  - Input Tokens:    {input_tokens} (Cost: {input_cost})\n"
                f"  - Thoughts Tokens: {thoughts_tokens} (Cost: {thoughts_cost})\n"
                f"  - Output Tokens:   {output_tokens} (Cost: {output_cost})\n"
                f"  - Total Tokens:    {total_tokens} (Cost: {total_cost})"
            )

            # Persist recent messages
            # await self.update_recent_messages(user_query, ai_msg)

            # Author-specific post-processing
            if author == AgentNameChoices.PRODUCT_RECOMMENDATION_AGENT.value:
                recommended_ids = agent_response.get('recommended_ids', [])
                source = agent_response.get('source', '')

                logger.info(f"Source of recommendation: {source}")
                logger.info(f"Recommended IDs: {recommended_ids}")

                if source and source.lower() == FeedSourceChoices.PRODUCT.value:
                    if tool_response_data and tool_response_name == 'combined_search':
                        # Tool already fetched products — reuse directly
                        recommended_results = tool_response_data.get('data', {}).get('recommended_results', {})
                    elif recommended_ids:
                        # Fetch from DB using recommended IDs
                        parent_product_ids = [
                            p.get("product_id") for p in recommended_ids if p.get("product_id")
                        ]
                        variant_ids = [
                            v for p in recommended_ids for v in (p.get("variant_ids") or [])
                        ]
                        products = get_products_by_ids(self.session_info.bot_id, parent_product_ids, variant_ids)
                        recommended_results = {
                            FeedSourceChoices.PRODUCT.value: products,
                            FeedSourceChoices.FAQ.value: []
                        }

                    if recommended_ids:
                        product_info = filter_products_by_recommendations(
                            vector_search_results=recommended_results,
                            recommendation_ids=recommended_ids
                        )
                        await self._persist_recommended_product_ids(recommended_ids)

                ai_response['product_info'] = product_info

            elif author == AgentNameChoices.ORDER_TRACKING_AGENT.value:
                order_info = tool_response_data.get("data", {}) if tool_response_data else {}
                has_asked_for_order_status = True
                ai_response['order_data'] = order_info

            elif author in (
                AgentNameChoices.GREETINGS_OR_FAREWELL_AGENT.value,
                AgentNameChoices.SMALL_TALK_AGENT.value,
                AgentNameChoices.HUMAN_ASSISTANT_NOTIFY_AGENT.value,
                AgentNameChoices.FAQ_RESPONSE_AGENT.value,
            ):
                pass  # no extra processing needed

            else:
                logger.info(f"Unhandled author: {author}")

            logger.info(f"Final AI response prepared for user: {ai_response}")

        except Exception as e:
            logger.error(f"Error occurred while processing agent response: {e}", exc_info=True)
            ai_response = {"ai_message": DEFAULT_ERROR_RESPONSE, "product_info": []}
            error = True

        return ChatResponseModel(
            session_id=self.session_info.user_session_id,
            ai_response=ChatResponseContentModel(
                ai_response=ai_response,
                token=total_tokens,
                cost=total_cost,
                error=error,
                notification=admin_notification,
                has_asked_for_order_status=has_asked_for_order_status,
                meta_data=ChatResponseMetaDataModel(
                    message_received_time_by_ai=self.session_info.sent_time,
                    query_intention=query_intention,
                    reply_to_message_id=self.session_info.reply_to_message_id
                )
            )
        )

代码的调用方式是：

client  = await AgentClient.create(session_info)

await client.chat()

在对该代码进行压力测试时，会出现以下错误：

The session has been modified in storage since it was loaded. Please reload the session before appending more events.

如何防止这个错误发生并让我的服务器在压力测试下正常运行？

解决方案

所以，我认为你遇到了一个经典的乐观并发控制（OCC）失败，这种失败是由跨异步边界的竞态条件引起的。在你的 build_search_query 方法中，你从数据库获取会话对象，等待通过 await llm.ask() 进行的慢速LLM生成，然后尝试使用那个原始会话对象执行 append_event。在压力测试下，对同一用户的并发请求会在那次LLM调用仍在挂起时就修改数据库中的会话。等到你的代码最终尝试附加新事件时，会话对象的内部版本哈希值已经完全过时，ADK存储引擎会抛出该修改错误以防止数据覆盖。

为了解决这个问题并在高负载下稳定服务器，你必须在每一次执行 append_event 之前，立即从数据库重新获取会话对象，而不是依赖于在函数开头加载的过时对象。你需要对 _persist_recommended_product_ids 和 update_recent_messages 应用同样的“重新加载再写入”的修复，以确保在推送增量更新之前会话状态始终是最新的。