Statnett-255: OBO authentication against Cognite

Author: nelly-hateva

CHANGELOG.md

@@ -1,10 +1,11 @@
 Next release
 ============
 
+* [#255](https://github.com/statnett/Talk2PowerSystem_PM/issues/255): OBO auth flow for Cognite
+* [#255](https://github.com/statnett/Talk2PowerSystem_PM/issues/255): Update the version of `cognite-sdk` from `7.86.0` to `7.89.0`
 * [#251](https://github.com/statnett/Talk2PowerSystem_PM/issues/251): Change N-Shot tool configuration to default to the base GraphDB
 * [#251](https://github.com/statnett/Talk2PowerSystem_PM/issues/251): Change N-Shot tool configuration SPARQL query template
 * [#254](https://github.com/statnett/Talk2PowerSystem_PM/issues/254): Update the version of `ttyg` from `1.9.3` to `1.10.0`, so that the chat bot can run without admin access to GraphDB
-* [#254](https://github.com/statnett/Talk2PowerSystem_PM/issues/254): Update the version of `cognite-sdk` from `7.86.0` to `7.88.0`
 
 1.2.0-rc4
 ============

docs/AgentConfig.md

@@ -150,6 +150,7 @@ One of `dev1`, `dev2`, `dev3`, `test`, `prod` according to [CDF access from RNDP
   Otherwise, `tools.cognite.token_file_path` must be provided for client credentials authentication.
 - `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.client_secret` - OPTIONAL - Client secret for the Cognite confidential application.
 
 ## `llm`
 

poetry.lock

@@ -13,7 +13,7 @@ dependencies = [
     "langchain-openai==0.3.32",
     "langgraph-checkpoint-redis==0.1.2",
     "jsonlines==4.0.0",
-    "cognite-sdk==7.88.0",
+    "cognite-sdk==7.89.0",
     "pydantic-settings==2.10.1",
     "PyYAML==6.0.2",
     "uvicorn[standard] (==0.35.0)",
@@ -22,6 +22,7 @@ dependencies = [
     "toml==0.10.2",
     "markdown==3.8.2",
     "python-jose[cryptography] (==3.5.0)",
+    "msal==1.34.0",
     "cachetools==6.2.0",
     "importlib_resources==6.5.2",
 ]

pyproject.toml

@@ -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",
-    "messages = {\"messages\": [(\"user\", \"List all transformers within substation OSLO.\")]}\n",
-    "last_message_id = run_agent(agent_executor, messages, conf)"
+    "conf = RunnableConfig(configurable={\"thread_id\": \"thread-123\"})\n",
+    "messages = {\"messages\": [(\"user\", \"List timeseries\")]}\n",
+    "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.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/jupyter_notebooks/Talk2PowerSystem_RNDP.ipynb

@@ -77,6 +77,7 @@ 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":
@@ -179,14 +180,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 +213,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 +228,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 +238,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 +263,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_react_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_react_agent(
             model=self.model,
-            tools=self.tools,
-            prompt=instructions,
-            checkpointer=checkpointer,
+            tools=tools,
+            prompt=self.instructions,
+            checkpointer=self.checkpointer,
         )

src/talk2powersystemllm/agent.py

@@ -58,7 +58,7 @@ def check_required_fields_and_set_default_oidc_discovery_url(self) -> "SecurityS
                              "client_id, authority, logout, login_redirect, logout_redirect, audience!")
         if self.enabled and not self.oidc_discovery_url:
             self.oidc_discovery_url = (
-                    self.authority.rstrip("/") + "/v2.0/.well-known/openid-configuration"
+                self.authority.rstrip("/") + "/v2.0/.well-known/openid-configuration"
             )
         return self
 
@@ -96,7 +96,7 @@ class AppSettings(BaseSettings):
     gtg_refresh_interval: int | None = Field(default=30, ge=1,
                                              description="The __gtg endpoint refresh interval in seconds")
     about_refresh_interval: int | None = Field(default=30, ge=1,
-                                             description="The __about endpoint refresh interval in seconds")
+                                               description="The __about endpoint refresh interval in seconds")
     trouble_md_path: Path | None = "/code/trouble.md"
     docs_url: str | None = "/docs"
     root_path: str | None = "/"

src/talk2powersystemllm/app/server/config.py

@@ -1,9 +1,7 @@
-from .cognite_healthcheck import CogniteHealthchecker
 from .graphdb_healthcheck import GraphDBHealthchecker
 from .healthchecks import HealthChecks
 
 __all__ = [
     "HealthChecks",
     "GraphDBHealthchecker",
-    "CogniteHealthchecker",
 ]