Statnett-255: OBO authentication against Cognite for the chat bot backend app

CHANGELOG.md

@@ -1,6 +1,7 @@
-Next release
+1.4.0-rc1
 ============
 
+* [#255](https://github.com/statnett/Talk2PowerSystem_PM/issues/255): OBO auth flow for Cognite
 * [#278](https://github.com/statnett/Talk2PowerSystem_PM/issues/278): Update the version of `graphrag-eval` from `5.2.0` to `5.3.1`
 
 1.3.0-rc1

docs/AgentConfig.md

@@ -150,10 +150,11 @@ LIMIT {limit}
 - `tools.cognite.project` - OPTIONAL, DEFAULT=`prod` - Cognite Data Fusion project name.
 One of `dev1`, `dev2`, `dev3`, `test`, `prod` according to [CDF access from RNDP](https://github.com/statnett/Talk2PowerSystem_PM/wiki/CDF-access-from-RNDP).
 - `tools.cognite.client_name` - OPTIONAL, DEFAULT=`talk2powersystem` - Name of the client for logging purposes.
-- `tools.cognite.interactive_client_id` - OPTIONAL - If provided, interactive authentication is used.
-  Otherwise, `tools.cognite.token_file_path` must be provided for client credentials authentication.
+- `tools.cognite.interactive_client_id` - OPTIONAL - If provided, interactive authentication is used (when you run on a dev machine the backend app with uvicorn or the Jupyter Notebook).
+  Otherwise, `tools.cognite.token_file_path` or `tools.cognite.client_secret` must be provided.
 - `tools.cognite.tenant_id` - REQUIRED iff `tools.cognite.interactive_client_id` is present - Azure tenant ID. For example, `a8d61462-f252-44b2-bf6a-d7231960c041`.
-- `tools.cognite.token_file_path` - OPTIONAL - Full path on the disk to the cognite token file. For example, `/var/run/secrets/microsoft.com/entra/cognite`.
+- `tools.cognite.token_file_path` - OPTIONAL - Full path on the disk to the cognite token file (used when you run the Jupyter Notebook on RNDP). For example, `/var/run/secrets/microsoft.com/entra/cognite`.
+* `tools.cognite.client_secret` - OPTIONAL - Client secret for the Cognite confidential application (used for the backend app running on RNDP).
 
 ## `llm`
 

docs/FastAPI.md

@@ -69,12 +69,14 @@ The `version` information is included in the response from the `/__about` endpoi
 
 * `SECURITY_ENABLED` - OPTIONAL, DEFAULT=False, Exposed to the UI - Indicates if security is enabled.
 * `SECURITY_CLIENT_ID` - REQUIRED iff `SECURITY_ENABLED=True`, Exposed to the UI - The registered application (client) ID.
+* `SECURITY_FRONTEND_APP_CLIENT_ID` - REQUIRED iff `SECURITY_ENABLED=True`, Exposed to the UI - The registered frontend application (client) ID.
 * `SECURITY_AUTHORITY` - REQUIRED iff `SECURITY_ENABLED=True`, Exposed to the UI - The authority URL used for authentication.
 * `SECURITY_LOGOUT` - REQUIRED iff `SECURITY_ENABLED=True`, Exposed to the UI - The logout endpoint URL.
 * `SECURITY_LOGIN_REDIRECT` - REQUIRED iff `SECURITY_ENABLED=True`, Exposed to the UI - The URL to redirect to after a successful login.
 * `SECURITY_LOGOUT_REDIRECT` - REQUIRED iff `SECURITY_ENABLED=True`, Exposed to the UI - The URL to redirect to after logout.
 * `SECURITY_OIDC_DISCOVERY_URL` - OPTIONAL, DEFAULT=`{SECURITY_AUTHORITY}/v2.0/.well-known/openid-configuration` - OpenID Connect Discovery URL.
 * `SECURITY_AUDIENCE` - REQUIRED iff `SECURITY_ENABLED=True` - The expected audience of the security tokens.
+* `SECURITY_ISSUER` - REQUIRED iff `SECURITY_ENABLED=True` - The expected issuer of the security tokens.
 * `SECURITY_TTL` - OPTIONAL, DEFAULT=`86400` seconds (24 hours), must be >= 1 - Indicates how many seconds to cache the public keys and the issuer obtained from the OpenID Configuration endpoint.
 According to [the Azure documentation](https://learn.microsoft.com/en-us/entra/identity-platform/access-tokens) a reasonable frequency to check for updates to the public keys used by Microsoft Entra ID is every 24 hours.
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Talk2PowerSystemLLM"
-version = "1.3.1-dev0"
+version = "1.4.0-rc1"
 description = "Talk to Power System LLM"
 authors = []
 readme = "README.adoc"
@@ -22,6 +22,7 @@ dependencies = [
     "toml==0.10.2",
     "markdown==3.10",
     "python-jose[cryptography] (==3.5.0)",
+    "msal==1.34.0",
     "cachetools==6.2.2",
     "importlib_resources==6.5.2",
 ]

src/jupyter_notebooks/Talk2PowerSystem.ipynb

@@ -79,12 +79,12 @@
     "\n",
     "from langgraph.checkpoint.memory import InMemorySaver\n",
     "\n",
-    "from talk2powersystemllm.agent import Talk2PowerSystemAgent\n",
+    "from talk2powersystemllm.agent import Talk2PowerSystemAgentFactory\n",
     "\n",
-    "agent_executor = Talk2PowerSystemAgent(\n",
+    "agent = Talk2PowerSystemAgentFactory(\n",
     "    Path(\"../../config/dev+cognite.yaml\"),\n",
     "    checkpointer=InMemorySaver()\n",
-    ").agent"
+    ").get_agent()"
    ]
   },
   {
@@ -106,12 +106,13 @@
    },
    "outputs": [],
    "source": [
+    "from langchain_core.runnables import RunnableConfig\n",
     "from ttyg.agents import run_agent\n",
     "\n",
     "\n",
-    "conf = {\"configurable\": {\"thread_id\": \"thread-123\"}}\n",
+    "conf = RunnableConfig(configurable={\"thread_id\": \"thread-123\"})\n",
     "messages = {\"messages\": [(\"user\", \"List all transformers within substation OSLO.\")]}\n",
-    "last_message_id = run_agent(agent_executor, messages, conf)"
+    "last_message_id = run_agent(agent, messages, conf)"
    ]
   },
   {
@@ -122,7 +123,7 @@
    "outputs": [],
    "source": [
     "messages = {\"messages\": [(\"user\", \"Give me their descriptions\")]}\n",
-    "last_message_id = run_agent(agent_executor, messages, conf, last_message_id=last_message_id)"
+    "last_message_id = run_agent(agent, messages, conf, last_message_id=last_message_id)"
    ]
   },
   {
@@ -144,6 +145,7 @@
    "source": [
     "from datetime import datetime\n",
     "\n",
+    "from langchain_core.runnables import RunnableConfig\n",
     "from ttyg.agents import run_agent\n",
     "\n",
     "from talk2powersystemllm.tools.user_datetime_context import user_datetime_ctx\n",
@@ -180,10 +182,10 @@
     "]\n",
     "\n",
     "for i, question in enumerate(questions):\n",
-    "    conf = {\"configurable\": {\"thread_id\": f\"thread-{i}\"}}\n",
+    "    conf = RunnableConfig(configurable={\"thread_id\": f\"thread-{i}\"})\n",
     "    messages = {\"messages\": [(\"user\", question)]}\n",
     "    user_datetime_ctx.set(datetime.now().astimezone().strftime(\"%Y-%m-%dT%H:%M:%S%z\"))\n",
-    "    run_agent(agent_executor, messages, conf)"
+    "    run_agent(agent, messages, conf)"
    ]
   }
  ],

src/jupyter_notebooks/Talk2PowerSystem_RNDP.ipynb

@@ -80,12 +80,12 @@
     "\n",
     "from langgraph.checkpoint.memory import InMemorySaver\n",
     "\n",
-    "from talk2powersystemllm.agent import Talk2PowerSystemAgent\n",
+    "from talk2powersystemllm.agent import Talk2PowerSystemAgentFactory\n",
     "\n",
-    "agent_executor = Talk2PowerSystemAgent(\n",
+    "agent = Talk2PowerSystemAgentFactory(\n",
     "    Path(\"../../config/rndp.yaml\"),\n",
     "    checkpointer=InMemorySaver()\n",
-    ").agent"
+    ").get_agent()"
    ]
   },
   {
@@ -107,12 +107,13 @@
    },
    "outputs": [],
    "source": [
+    "from langchain_core.runnables import RunnableConfig\n",
     "from ttyg.agents import run_agent\n",
     "\n",
     "\n",
-    "conf = {\"configurable\": {\"thread_id\": \"thread-123\"}}\n",
+    "conf = RunnableConfig(configurable={\"thread_id\": \"thread-123\"})\n",
     "messages = {\"messages\": [(\"user\", \"List all transformers within substation OSLO.\")]}\n",
-    "last_message_id = run_agent(agent_executor, messages, conf)"
+    "last_message_id = run_agent(agent, messages, conf)"
    ]
   },
   {
@@ -123,7 +124,7 @@
    "outputs": [],
    "source": [
     "messages = {\"messages\": [(\"user\", \"Give me their descriptions\")]}\n",
-    "last_message_id = run_agent(agent_executor, messages, conf, last_message_id=last_message_id)"
+    "last_message_id = run_agent(agent, messages, conf, last_message_id=last_message_id)"
    ]
   },
   {
@@ -145,6 +146,7 @@
    "source": [
     "from datetime import datetime\n",
     "\n",
+    "from langchain_core.runnables import RunnableConfig\n",
     "from ttyg.agents import run_agent\n",
     "\n",
     "from talk2powersystemllm.tools.user_datetime_context import user_datetime_ctx\n",
@@ -181,10 +183,10 @@
     "]\n",
     "\n",
     "for i, question in enumerate(questions):\n",
-    "    conf = {\"configurable\": {\"thread_id\": f\"thread-{i}\"}}\n",
+    "    conf = RunnableConfig(configurable={\"thread_id\": f\"thread-{i}\"})\n",
     "    messages = {\"messages\": [(\"user\", question)]}\n",
     "    user_datetime_ctx.set(datetime.now().astimezone().strftime(\"%Y-%m-%dT%H:%M:%S%z\"))\n",
-    "    run_agent(agent_executor, messages, conf)"
+    "    run_agent(agent, messages, conf)"
    ]
   }
  ],

src/talk2powersystemllm/agent.py

@@ -1,6 +1,7 @@
 from base64 import b64encode
 from enum import Enum
 from pathlib import Path
+from typing import Any
 
 import yaml
 from langchain.agents import create_agent
@@ -77,12 +78,16 @@ class CogniteSettings(BaseModel):
     token_file_path: Path | None = None
     interactive_client_id: str | None = None
     tenant_id: str | None = None
+    client_secret: str | None = None
 
     @model_validator(mode="after")
     def check_credentials(self) -> "CogniteSettings":
-        if self.token_file_path and self.interactive_client_id:
-            raise ValueError("Both token_file_path and interactive_client_id for Cognite are provided. "
-                             "Set only one of them!")
+        def exactly_one_is_not_none(*args: Any) -> bool:
+            return sum(a is not None for a in args) == 1
+
+        if not exactly_one_is_not_none(self.client_secret, self.token_file_path, self.interactive_client_id):
+            raise ValueError("Pass exactly one of 'client_secret', 'token_file_path' or 'interactive_client_id'!")
+
         if self.interactive_client_id and not self.tenant_id:
             raise ValueError("Tenant id is required!")
         return self
@@ -179,14 +184,15 @@ def init_graphdb(graphdb_settings: GraphDBSettings) -> GraphDB:
     return GraphDB(**kwargs)
 
 
-def init_cognite(cognite_settings: CogniteSettings) -> CogniteSession:
+def init_cognite(cognite_settings: CogniteSettings, obo_token: str | None = None) -> CogniteSession:
     return CogniteSession(
         base_url=cognite_settings.base_url,
         client_name=cognite_settings.client_name,
         project=cognite_settings.project,
         token_file_path=cognite_settings.token_file_path,
         interactive_client_id=cognite_settings.interactive_client_id,
         tenant_id=cognite_settings.tenant_id,
+        obo_token=obo_token,
     )
 
 
@@ -211,11 +217,13 @@ def init_llm(llm_settings: LLMSettings) -> BaseChatModel:
         )
 
 
-class Talk2PowerSystemAgent:
-    agent: CompiledStateGraph
+class Talk2PowerSystemAgentFactory:
+    settings: Talk2PowerSystemAgentSettings
     graphdb_client: GraphDB
-    cognite_session: CogniteSession | None = None
+    checkpointer: Checkpointer | None = None
     model: BaseChatModel
+    instructions: str
+    tools: list[BaseTool]
 
     def __init__(
         self,
@@ -224,6 +232,7 @@ def __init__(
     ):
         self.settings = read_config(path_to_yaml_config)
         self.graphdb_client = init_graphdb(self.settings.graphdb)
+        self.checkpointer = checkpointer
 
         tools_settings = self.settings.tools
         self.tools: list[BaseTool] = []
@@ -233,11 +242,6 @@ def __init__(
         )
         self.tools.append(sparql_query_tool)
 
-        ontology_schema_and_vocabulary_tool = OntologySchemaAndVocabularyTool(
-            graph=self.graphdb_client,
-            ontology_schema_file_path=tools_settings.ontology_schema.file_path,
-        )
-
         autocomplete_search_settings = tools_settings.autocomplete_search
         autocomplete_search_kwargs = {
             "property_path": autocomplete_search_settings.property_path,
@@ -263,22 +267,28 @@ def __init__(
             )
             self.tools.append(retrieval_query_tool)
 
-        if tools_settings.cognite:
-            cognite_settings = tools_settings.cognite
-            self.cognite_session = init_cognite(cognite_settings)
-            self.tools.append(RetrieveTimeSeriesTool(cognite_session=self.cognite_session))
-            self.tools.append(RetrieveDataPointsTool(cognite_session=self.cognite_session))
-
         self.tools.append(NowTool())
 
-        instructions = f"""{self.settings.prompts.assistant_instructions}""".replace(
+        ontology_schema_and_vocabulary_tool = OntologySchemaAndVocabularyTool(
+            graph=self.graphdb_client,
+            ontology_schema_file_path=tools_settings.ontology_schema.file_path,
+        )
+        self.instructions = f"""{self.settings.prompts.assistant_instructions}""".replace(
             "{ontology_schema}", ontology_schema_and_vocabulary_tool.schema_graph.serialize(format="turtle")
         )
 
         self.model = init_llm(self.settings.llm)
-        self.agent = create_agent(
+
+    def get_agent(self, cognite_obo_token: str | None = None) -> CompiledStateGraph:
+        tools = self.tools.copy()
+        if self.settings.tools.cognite:
+            cognite_settings = self.settings.tools.cognite
+            cognite_session = init_cognite(cognite_settings, cognite_obo_token)
+            tools.append(RetrieveTimeSeriesTool(cognite_session=cognite_session))
+            tools.append(RetrieveDataPointsTool(cognite_session=cognite_session))
+        return create_agent(
             model=self.model,
-            tools=self.tools,
-            system_prompt=instructions,
-            checkpointer=checkpointer,
+            tools=tools,
+            system_prompt=self.instructions,
+            checkpointer=self.checkpointer,
         )

src/talk2powersystemllm/app/models/auth.py

@@ -4,6 +4,8 @@
 class AuthConfig(BaseModel):
     enabled: bool
     client_id: str | None = Field(alias="clientId")
+    frontend_app_client_id: str | None = Field(alias="frontendAppClientId")
+    scopes: list[str] | None
     authority: str | None
     logout: str | None
     login_redirect: str | None = Field(alias="loginRedirect")