File size: 8,005 Bytes
c3674d7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
240735d
 
 
 
 
 
 
 
 
c3674d7
240735d
 
 
 
 
c3674d7
 
 
240735d
 
c3674d7
 
 
 
240735d
 
c3674d7
 
240735d
 
 
c3674d7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
240735d
 
 
 
c3674d7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
"""
Clarification Question Handler
Manages the generation and collection of clarifying questions.
Backend functions are called directly (in-process) rather than via HTTP to avoid
self-call reliability issues in the single-process FastAPI+Chainlit setup.
"""
from typing import Dict, Any, List, Optional
import asyncio
import os


class ClarificationState:
    """Tracks the state of clarifying questions for a session"""

    def __init__(self, questions: List[Dict[str, Any]], original_query: str):
        self.original_query = original_query
        self.questions = questions
        self.answers = {}
        self.current_index = 0

    def get_current_question(self) -> Optional[Dict[str, Any]]:
        if self.current_index < len(self.questions):
            return self.questions[self.current_index]
        return None

    def add_answer(self, question_id: int, answer: str) -> bool:
        current_question = self.get_current_question()
        if not current_question or current_question["id"] != question_id:
            return False
        question_id_str = str(question_id)
        current_question["answer"] = answer
        self.answers[question_id_str] = {
            "question": current_question["question"],
            "answer": answer,
        }
        self.current_index += 1
        return True

    def is_complete(self) -> bool:
        return self.current_index >= len(self.questions)

    def get_progress(self) -> Dict[str, Any]:
        return {
            "answered": len(self.answers),
            "total": len(self.questions),
            "current_index": self.current_index,
            "percentage": int((len(self.answers) / len(self.questions)) * 100) if self.questions else 100,
        }

    def get_summary(self) -> Dict[str, Any]:
        return {
            "original_query": self.original_query,
            "total_questions": len(self.questions),
            "answers": self.answers,
            "complete": self.is_complete(),
        }

    def to_dict(self) -> Dict[str, Any]:
        return {
            "original_query": self.original_query,
            "questions": self.questions,
            "answers": {str(k): v for k, v in self.answers.items()},
            "current_index": self.current_index,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ClarificationState":
        state = cls(
            questions=data.get("questions", []),
            original_query=data.get("original_query", ""),
        )
        state.answers = data.get("answers", {})
        state.current_index = data.get("current_index", 0)
        return state


_gemini_key_lock = asyncio.Lock()


_VERTEXAI_ENV_VARS = (
    "GOOGLE_GENAI_USE_VERTEXAI",
    "GOOGLE_GENAI_USE_ENTERPRISE",
    "GOOGLE_CLOUD_PROJECT",
    "GOOGLE_CLOUD_LOCATION",
    "GOOGLE_APPLICATION_CREDENTIALS",
)


async def _with_gemini_key(coro, api_key: Optional[str]):
    """Run *coro* using a user-supplied Gemini API key.

    Temporarily sets GOOGLE_API_KEY and clears any Vertex AI env vars so that
    google-genai uses Gemini API key auth instead of Vertex AI credentials.
    """
    if not api_key:
        return await coro
    async with _gemini_key_lock:
        orig_api_key = os.environ.get("GOOGLE_API_KEY")
        orig_vertexai = {k: os.environ.pop(k, None) for k in _VERTEXAI_ENV_VARS}
        os.environ["GOOGLE_API_KEY"] = api_key
        try:
            return await coro
        finally:
            if orig_api_key is not None:
                os.environ["GOOGLE_API_KEY"] = orig_api_key
            else:
                os.environ.pop("GOOGLE_API_KEY", None)
            for k, v in orig_vertexai.items():
                if v is not None:
                    os.environ[k] = v


class ClarificationHandler:
    """
    Handles the flow of clarifying questions.
    Calls backend functions directly (in-process) to avoid HTTP self-call issues.
    """

    def __init__(self, backend_url_resolver=None):
        self.backend_url_resolver = backend_url_resolver  # kept for API compat, unused

    async def generate_questions(
        self,
        query: str,
        model_provider: str = "gemma",
        api_key: Optional[str] = None,
    ) -> Optional[ClarificationState]:
        """Generate clarifying questions by calling the pipeline function directly."""
        try:
            if model_provider == "gemma":
                from backend.llm.gemma_pipeline import generate_cq
                result = await generate_cq(query)
            else:
                from backend.adk.agents.clar_q_gen.cq_generator import generate_clarifying_questions
                result = await _with_gemini_key(
                    generate_clarifying_questions(query), api_key
                )

            if result is None:
                print(f"[ClarificationHandler] No result from CQ generator for: {query}")
                return None

            # result may be a CQOutput pydantic model or a dict
            if hasattr(result, "clarifying_questions"):
                questions = [q.model_dump() for q in result.clarifying_questions]
                original_query = result.query
            else:
                questions = result.get("clarifying_questions", [])
                original_query = result.get("query", query)

            if not questions or (len(questions) == 1 and questions[0].get("id") == -1):
                print(f"[ClarificationHandler] Query is out of scope: {query}")
                return None

            return ClarificationState(questions=questions, original_query=original_query)

        except Exception as e:
            import traceback
            print(f"[ClarificationHandler] Error generating questions: {e}")
            traceback.print_exc()
            # Re-raise API/auth errors so the caller can show the user a real error
            # message rather than silently treating it as out-of-scope.
            if model_provider == "gemini":
                raise
            return None

    def format_question_for_ui(self, state: ClarificationState) -> Dict[str, Any]:
        current_question = state.get_current_question()
        progress = state.get_progress()

        if not current_question:
            return self._format_completion(state)

        return {
            "type": "clarification_question",
            "text": current_question["question"],
            "question_id": current_question["id"],
            "progress": progress,
            "metadata": {"clarification_active": True},
        }

    def _format_completion(self, state: ClarificationState) -> Dict[str, Any]:
        return {
            "type": "clarification_complete",
            "text": (
                "✅ **All questions answered!**\n\n"
                "Thank you for providing that information. "
                "Let me now find the best recommendations for you based on your preferences."
            ),
            "summary": state.get_summary(),
            "metadata": {
                "clarification_active": False,
                "clarification_complete": True,
            },
        }

    def format_error(self, error_message: str) -> Dict[str, Any]:
        return {
            "type": "clarification_error",
            "text": "⚠️ Unable to generate clarifying questions: Out of scope",
            "metadata": {"clarification_active": False, "error": True},
        }

    def format_out_of_scope(self) -> Dict[str, Any]:
        return {
            "type": "out_of_scope",
            "text": (
                "I'm sorry, but this query is beyond the scope of European city recommendation. "
                "I can only help you find and recommend European cities based on your preferences.\n\n"
                "Please ask a new query about European cities to start again."
            ),
            "metadata": {
                "clarification_active": False,
                "out_of_scope": True,
                "reset_session": True,
            },
        }