rewriter.py

rewriter.py

  The rewriter is a single LLM node that turns a follow-up into a standalone Bangla question. It is used in two
  places: the standalone POST /api/v1/rewrite/ endpoint, and as a node inside the chatbot graph that
  /api/v1/qa/answer runs first.

  Layout

  ec-llm-service/
    src/api/v1/rewrite/          HTTP wrapper
      schema.py                  request/response
      routes.py                  POST /
    src/genai/rewrite/           the actual rewriter
      state.py                   TypedDict
      prompt.py                  REWRITE_PROMPT
      graph.py                   rewrite_node()

  graph.py is not a compiled LangGraph. It is one async function that the chatbot graph later mounts as a node.

  State

   ec-llm-service/src/genai/rewrite/state.py lines 1-10

  from typing import TypedDict
  class RewriteState(TypedDict):
      """Rewrite a follow-up question into a standalone one using history."""
      question: str
      history: list[tuple[str, str]]
      rewritten_question: str

  Input is the current question plus (question, answer) turns. Output is rewritten_question.

  How rewrite_node works

   ec-llm-service/src/genai/rewrite/graph.py lines 14-41

  async def rewrite_node(state: RewriteState) -> RewriteState:
      """Rewrite a follow-up into a standalone question using session history.
      Falls back to the original question when there is no history or the
      rewrite fails — the pipeline must keep working without memory.
      """
      question = state["question"]
      history = state.get("history") or []
      if not history:
          state["rewritten_question"] = question
          return state
      history_text = "\n".join(f"ব্যবহারকারী: {q}\nবট: {a}" for q, a in history)
      prompt = REWRITE_PROMPT.format(history=history_text, question=question)
      try:
          async with llm_semaphore:
              result = await asyncio.wait_for(
                  model.ainvoke([HumanMessage(content=prompt)]),
                  timeout=settings.LLM_TIMEOUT_S,
              )
          rewritten = result.content.strip()
      except Exception as e:
          logger.warning("Question rewrite failed, using original: %s", e)
          state["rewritten_question"] = question
          return state
      # ... empty-text fallback, then write rewritten_question

  The flow is:

  1. No history → pass-through. Nothing to resolve, so the original question is kept.
  2. Format history as Bangla turns: ব্যবহারকারী: … / বট: ….
  3. Fill REWRITE_PROMPT with that history and the follow-up.
  4. Call the same local LLM used everywhere else (src/genai/llm.py), behind llm_semaphore and LLM_TIMEOUT_S.
  5. Fail open. Timeout, LLM error, or empty text → original question. Retrieval still runs.

  Unlike the router and selector, this is a free-text call. No structured output / tool calling. Temperature is 0.0.

  Prompt rules

  The prompt asks for one standalone Bangla question and nothing else. The important rules:

  • If the follow-up already names a subject, that is a new topic. Complete the sentence about that subject. Do
    not import an entity from earlier turns.
  • Only pull a subject from history when the follow-up has no subject (pronoun, “how long?”, “what document?”).
  • Never merge two different entities.
  • If the follow-up is already self-contained, return it unchanged.
  • Return only the question — no quotes, no explanation.

  Example from the API schema: history is “ভোটার আইডি কার্ড কিভাবে ये पाবো”, follow-up is “কতদিন সময় লাগবে?” → “ভোটার আইডি কার্ড পেতে কতদিন
  সময় লাগবে?”

  Standalone API

  Mounted in main.py as /api/v1/rewrite:

   ec-llm-service/src/api/v1/rewrite/routes.py lines 13-26

  @router.post("/", response_model=RewriteResponse)
  async def rewrite_query(payload: RewriteRequest) -> RewriteResponse:
      started_at = perf_counter()
      history = [(turn.question, turn.answer) for turn in payload.history]
      state = await rewrite_node({"question": payload.question, "history": history})
      rewritten = state["rewritten_question"]
      return RewriteResponse(
          question=payload.question,
          rewritten_question=rewritten,
          was_rewritten=rewritten != payload.question,
          response_time_s=round(perf_counter() - started_at, 3),
      )

  You supply history in the body. The route does not use session_memory. was_rewritten is a string compare of
  original vs rewritten.

  curl -X POST http://localhost:8000/api/v1/rewrite/ \
    -H "Content-Type: application/json" \
    -d '{
          "question": "কতদিন সময় লাগবে?",
          "history": [{"question": "ভোটার আইডি কার্ড কিভাবে পাবো",
                       "answer": "নিকটস্থ নির্বাচন অফিসে আবেদন করুন।"}]
        }'

  Same node in the QA pipeline

  The chatbot graph reuses rewrite_node as-is. ChatState is a superset of RewriteState (adds route / direct_answer),
  so the node can write rewritten_question on chatbot state.

   ec-llm-service/src/genai/chatbot/graph.py lines 55-70

  def build_graph():
      graph = StateGraph(ChatState)
      graph.add_node("rewrite", rewrite_node)
      graph.add_edge("rewrite", END)
      if settings.ENABLE_ROUTER_LLM:
          graph.add_node("router", router_node)
          graph.add_edge(START, "router")
          graph.add_conditional_edges(
              "router",
              route_after_router,
              {"end": END, "rewrite": "rewrite"},
          )
      else:
          graph.add_edge(START, "rewrite")

  /api/v1/qa/answer then:

  1. Loads history from in-memory session_memory (last 5 turns, LRU sessions).
  2. Runs chatbot_graph. General chit-chat skips rewrite and returns direct_answer.
  3. For EC queries, takes route_state["rewritten_question"] and sends that to retrieval/selector.
  4. Stores the original user question + bot answer in session memory, not the rewritten form.

   ec-llm-service/src/api/v1/qa/routes.py lines 28-35

      history = session_memory.get(session_id)
      route_state = await chatbot_graph.ainvoke(
          {"question": payload.question, "history": history}
      )
      # ... general route returns early ...
      question = route_state["rewritten_question"]
      response = await chatbot.answer(question, payload.top_k)
      session_memory.add(session_id, payload.question, response.answer)

  So the rewriter exists so retrieval sees a complete question (“ভোটার আইডি কার্ড পেতে কতদিন…”) instead of a dangling fol  So the rewriter exists so retrieval sees a complete question (“ভোটার আইডি কার্ড পেতে কতদিন…”) instead of a dangling
  follow-up (“কতদিন সময় লাগবে?”), while memory stays in the user’s original wording.