# # streaming_rag_agent.py
# import time
# from typing import List, Dict, AsyncGenerator, Optional

# import google.generativeai as genai

# class StreamingRAGAgent:
#     """Streaming RAG agent for tender documents."""

#     def __init__(self, collection, embed_fn):
#         self.collection = collection
#         self.embed_text = embed_fn
#         self.model = genai.GenerativeModel(
#             model_name="gemini-2.0-flash",
#             system_instruction="""
# <system>
# You are a smart call center AI assistant specialized in tender documents.
# - Answer concisely in a single short paragraph.
# - Use only retrieved content; do not speculate.
# - Format responses in plain text suitable for speech synthesis or agent scripts.
# - Avoid greetings, sign-offs, or unnecessary words.
# </system>
# """
#         )

#     def retrieve(self, query: str, top_k: int = 2, kb_ids: Optional[List[str]] = None) -> List[Dict]:
#         """Retrieve top-k relevant documents from Milvus."""
#         filter_expr = None
#         if kb_ids:
#             # Assuming documents have a 'kb_id' field in Milvus
#             filter_expr = f'knowledge_base_ids in {kb_ids}'
            
#         q_vec = self.embed_text(query)
#         results = self.collection.search(
#             data=[q_vec],
#             anns_field="embedding",
#             param={"metric_type": "IP", "params": {"nprobe": 10}},
#             limit=top_k,
#             output_fields=["window"],
#             expr = filter_expr
#         )
#         return [{"text": hit.entity.get("window"), "distance": hit.distance} for hit in results[0]]

#     async def stream_answer(self, user_query: str, memory: Optional[List[str]] = None, kb_ids: Optional[List[str]] = None) -> AsyncGenerator[str, None]:
#         """Async generator yielding tokens from the model."""
#         if memory is None:
#             memory = []
#         try:
#             retrieved = self.retrieve(user_query,2,kb_ids)
#         except Exception as e:
#             yield f"<error>Retrieval failed: {e}</error>"
#             return

#         prompt = f"""
# <user_query>{user_query}</user_query>
# <retrieved>{retrieved}</retrieved>
# <previous_user_queries>{memory}</previous_user_queries>
# <instruction>Answer concisely using only retrieved content.</instruction>
# """

#         start_time = time.time()
#         first_token_time = None

#         try:
#             stream = self.model.generate_content(
#                 contents=[{"role": "user", "parts": [{"text": prompt}]}],
#                 stream=True
#             )
#             for chunk in stream:
#                 if not chunk.candidates:
#                     continue
#                 token = "".join(p.text for p in chunk.candidates[0].content.parts)
#                 if token.strip() and first_token_time is None:
#                     first_token_time = time.time() - start_time
#                     print(f"\n<timing>Time to first token: {first_token_time:.3f}s</timing>\n")
#                 yield token
#         except Exception as e:
#             yield f"<error>Generation failed: {e}</error>"