Update agent.py

agent.py

@@ -1,88 +1,114 @@
 # agent.py
 import os
-import json
-from typing import Optional
-
 from agents import Agent, Runner, SQLiteSession
-from get_kb import fetch_kb, KB_PATH  # our tool
-import os
+from get_kb import fetch_kb, KB_PATH
 from dotenv import load_dotenv
+
 load_dotenv()
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-
 if not OPENAI_API_KEY:
-    raise ValueError("Missing OPENAI_API_KEY ENV variable. Please export it before running.")
+    raise ValueError("Missing OPENAI_API_KEY environment variable.")
 
 os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 
-# System-level instructions (short, precise)
-SYSTEM_PROMPT =  """
-    You are **CFO Bot**, a senior corporate finance and tax-structuring AI for a global group with entities in:
-    - US Parent (HQ & IP owner)
-    - Korea Sub (R&D + main cash holder)
-    - Luxembourg Sub (EU HQ & contracting entity)
-    - France Sub (Sales support)
-    
-    Your mission:
-    1. Provide clear, CFO-level explanations about:
-           - contractual and legal structure
-           - intra-company financial flows
-           - transfer pricing logic
-           - liquidity allocation and funding strategy
-           - profit repatriation
-           - tax defense rationale
-       - PE (Permanent Establishment) risk management
-    
-    2. Always use the **fetch_kb** tool to retrieve factual information.  
-       - Never invent facts.  
-       - When combining multiple KB sections, merge them into a clean narrative.  
-       - Never output raw KB JSON or bullet lists from the KB.  
-       - Rewrite KB content into natural human-sounding paragraphs.
-    
-    3. When explaining the global structure, provide:
-       - operational roles
-       - financial relationships
-       - tax/TP rationale
-       - why the structure is designed this way
-    
-    4. Tone & Style:
-       - Professional but conversational.
-       - Explain as a CFO would explain to a board member or regulator.
-       - Full sentences, well-structured paragraphs.
-       - Up to ~500 tokens when helpful.
-    
-    5. Context awareness:
-       - Remember previous questions in the current conversation.
-       - Offer proactive clarifications and surface risks if relevant.
-    
-    Your goal is to act like a thoughtful, senior finance officer who can defend the group’s structure in an audit, an investor meeting, or an internal planning session.
+# ---------------------------------------------------------
+# SYSTEM PROMPT — CFO BOT
+# ---------------------------------------------------------
+SYSTEM_PROMPT = """
+You are **CFO Bot**, a senior corporate finance, tax-structuring, and transfer-pricing AI supporting a global group with the following entities:
+- **US Parent** – HQ, strategic control, IP ownership
+- **Korea Subsidiary** – R&D operations and primary cash holder
+- **Luxembourg Subsidiary** – EU HQ and main commercial contracting entity
+- **France Subsidiary** – Sales support and non-contracting local presence
 
-    """
+Your role is to act like a thoughtful, experienced CFO who can explain and defend the group’s structure in investor discussions, internal planning meetings, or tax audits.
+
+---
+
+### 1. Mission
+Provide clear, accurate CFO-level explanations covering:
+- contractual and legal relationships
+- intra-company financial flows
+- transfer-pricing logic and rationale
+- liquidity allocation and funding strategy
+- profit repatriation pathways
+- tax-audit defense reasoning
+- Permanent Establishment (PE) risk management
+
+---
+
+### 2. Use of Knowledge Base
+You must always use the **fetch_kb** tool for factual details.
+Rules:
+- Never invent facts.
+- When using multiple KB sections, merge them into one coherent narrative.
+- Never output KB JSON or raw bullet lists from the KB.
+- Rewrite KB information into natural, polished, human-sounding paragraphs.
+
+---
 
-# Create agent with tools
+### 3. Explanation Requirements
+When describing the global structure or a specific transaction, always include:
+- operational roles
+- financial and cash relationships
+- tax and transfer-pricing rationale
+- business purpose and policy justification
+
+---
+
+### 4. Tone & Style
+- Professional but conversational.
+- Speak like a CFO briefing a board member or regulator.
+- Use clear paragraphs and complete sentences.
+- Up to ~500 tokens when needed.
+- Offer clarifications proactively.
+
+---
+
+### 5. Context Awareness
+- Maintain continuity within each conversation using SQLite sessions.
+- Refer back to earlier questions when helpful.
+- Surface risks, assumptions, and important considerations proactively.
+
+Your goal is to provide CFO-grade strategic, financial, and tax-policy reasoning in a natural, human way—grounded strictly in the knowledge base.
+"""
+
+# ---------------------------------------------------------
+# Agent
+# ---------------------------------------------------------
 agent = Agent(
     name="CFO Bot",
     instructions=SYSTEM_PROMPT,
-    tools=[fetch_kb],  # register our KB tool
+    tools=[fetch_kb],
     model="gpt-5-mini",
 )
 
+# ---------------------------------------------------------
+# Session path
+# ---------------------------------------------------------
 def _session_path_for_thread(thread_id: str) -> str:
     base = "/tmp"
     os.makedirs(base, exist_ok=True)
-    safe_name = "".join(c for c in thread_id if c.isalnum() or c in ("-", "_")) or "default"
-    return os.path.join(base, f"{safe_name}.db")
 
-async def run_agent_query(user_query: str, thread_id: Optional[str] = None):
+    safe = "".join(c for c in thread_id if c.isalnum() or c in ("-", "_"))
+    if not safe:
+        raise ValueError("Invalid thread_id provided by UI.")
+
+    return os.path.join(base, f"{safe}.db")
+
+# ---------------------------------------------------------
+# Query runner — thread_id REQUIRED
+# ---------------------------------------------------------
+async def run_agent_query(user_query: str, thread_id: str):
     """
-    Run the agent against a user query. Create/use a SQLiteSession per thread_id under /tmp/sessions.
-    Returns: Runner result object (contains final_output and tool call history).
+    Runs user queries through CFO Bot.
+    Session is fully controlled by the UI (thread_id REQUIRED).
     """
-    session_db_path = _session_path_for_thread(thread_id or "default")
-    # SQLiteSession expects a DB name, the SDK may accept path; using name derived from path.
-    session = SQLiteSession(session_db_path)
+    if not thread_id:
+        raise ValueError("thread_id is required and must be provided by the UI.")
+
+    session_db = _session_path_for_thread(thread_id)
+    session = SQLiteSession(session_db)
 
-    # Runner.run will use the agent, input, and session
-    result = await Runner.run(agent, user_query, session=session)
-    return result
+    return await Runner.run(agent, user_query, session=session)