Add ChatOpenRouter class and integrate system prompt loading in BasicAgent

agent.py

@@ -3,7 +3,8 @@ from typing import Optional
 
 from langgraph.graph import StateGraph, START, END, MessagesState
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
-from langgraph.prebuilt import ToolNode  # add
+from langgraph.prebuilt import ToolNode
+from pydantic import Field, SecretStr  # add
 
 # Try to import tools from tools.py
 try:
@@ -21,6 +22,29 @@ except Exception:  # pragma: no cover - optional dependency resolution
     ChatOpenAI = None  # type: ignore
 
 
+class ChatOpenRouter(ChatOpenAI):
+    openai_api_key: Optional[SecretStr] = Field(
+        alias="api_key",
+        default_factory=os.getenv("OPENROUTER_API_KEY", None),
+    )
+
+    @property
+    def lc_secrets(self) -> dict[str, str]:
+        return {"openai_api_key": "OPENROUTER_API_KEY"}
+
+    def __init__(self,
+                 openai_api_key: Optional[str] = None,
+                 **kwargs):
+        openai_api_key = (
+            openai_api_key or os.getenv("OPENROUTER_API_KEY")
+        )
+        super().__init__(
+            base_url="https://openrouter.ai/api/v1",
+            openai_api_key=openai_api_key,
+            **kwargs
+        )
+
+
 class _EchoModel:
     """Simple stub model used when no API key / model is configured."""
 
@@ -56,14 +80,15 @@ class LangGraphAgent:
 
         # Choose an LLM if not provided
         if model is None:
-            if ChatOpenAI is not None:
+            model = ChatOpenRouter(model="openai/gpt-oss-20b:free")
+            if model is None and ChatOpenAI is not None:
                 model = ChatOpenAI(
                     api_key=os.getenv("OPENROUTER_API_KEY"),
                     base_url=os.getenv("OPENROUTER_BASE_URL"),
                     model="openai/gpt-oss-20b:free",
                 )
-            else:
-                model = _EchoModel()
+        if model is None:
+            model = _EchoModel()
         self.model = model
 
         # Load tools and bind to the model if supported

app.py

@@ -4,6 +4,32 @@ import requests
 import inspect
 import pandas as pd
 
+# --- System Prompt Integration ---
+try:
+    from prompt import get_system_prompt as _get_system_prompt  # preferred
+except Exception:
+    _get_system_prompt = None  # type: ignore
+try:
+    from prompt import SYSTEM_PROMPT as _SYSTEM_PROMPT  # fallback constant
+except Exception:
+    _SYSTEM_PROMPT = None  # type: ignore
+
+
+def _load_system_prompt() -> str:
+    """
+    Resolve system prompt from prompt.py:
+    - Prefer get_system_prompt() if available, else SYSTEM_PROMPT, else default.
+    """
+    if callable(_get_system_prompt):
+        try:
+            return _get_system_prompt()  # type: ignore
+        except Exception:
+            pass
+    if isinstance(_SYSTEM_PROMPT, str) and _SYSTEM_PROMPT.strip():
+        return _SYSTEM_PROMPT  # type: ignore
+    return "You are a helpful assistant. Keep answers concise."
+
+
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -16,8 +42,9 @@ class BasicAgent:
     def __init__(self):
         from agent import LangGraphAgent
         print("BasicAgent initialized (LangGraph).")
-        # Create a minimal LangGraph agent; will use OPENAI_API_KEY if set, else a stub echo
-        self._agent = LangGraphAgent()
+        # Create a LangGraph agent with system prompt from prompt.py
+        system_prompt = _load_system_prompt()
+        self._agent = LangGraphAgent(system_prompt=system_prompt)
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")