Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 27 additions & 24 deletions cognite/client/_api/agents/agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,10 +54,9 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... AgentUpsert,
... QueryKnowledgeGraphAgentToolUpsert,
... QueryKnowledgeGraphAgentToolConfiguration,
... DataModelInfo
... DataModelInfo,
... )
>>> client = CogniteClient()
...
>>> find_assets_tool = QueryKnowledgeGraphAgentToolUpsert(
... name="find assets",
... description="Use this tool to find assets",
Expand All @@ -70,13 +69,13 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... view_external_ids=["CogniteAsset"],
... )
... ]
... )
... ),
... )
>>> agent = AgentUpsert(
... external_id="my_agent",
... name="My Agent",
... labels=["published"],
... tools=[find_assets_tool]
... tools=[find_assets_tool],
... )
>>> client.agents.upsert(agents=[agent])

Expand All @@ -89,9 +88,8 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... DataModelInfo,
... SummarizeDocumentAgentToolUpsert,
... AskDocumentAgentToolUpsert,
... QueryTimeSeriesDatapointsAgentToolUpsert
... QueryTimeSeriesDatapointsAgentToolUpsert,
... )
...
>>> find_assets_tool = QueryKnowledgeGraphAgentToolUpsert(
... name="find assets",
... description="Use this tool to query the knowledge graph for assets",
Expand All @@ -104,7 +102,7 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... view_external_ids=["CogniteAsset"],
... )
... ]
... )
... ),
... )
>>> find_files_tool = QueryKnowledgeGraphAgentToolUpsert(
... name="find files",
Expand All @@ -118,7 +116,7 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... view_external_ids=["CogniteFile"],
... )
... ]
... )
... ),
... )
>>> find_time_series_tool = QueryKnowledgeGraphAgentToolUpsert(
... name="find time series",
Expand All @@ -132,27 +130,34 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... view_external_ids=["CogniteTimeSeries"],
... )
... ]
... )
... ),
... )
>>> summarize_tool = SummarizeDocumentAgentToolUpsert(
... name="summarize document",
... description="Use this tool to get a summary of a document"
... description="Use this tool to get a summary of a document",
... )
>>> ask_doc_tool = AskDocumentAgentToolUpsert(
... name="ask document",
... description="Use this tool to ask questions about specific documents"
... description="Use this tool to ask questions about specific documents",
... )
>>> ts_tool = QueryTimeSeriesDatapointsAgentToolUpsert(
... name="query time series",
... description="Use this tool to query time series data points"
... description="Use this tool to query time series data points",
... )
>>> agent = AgentUpsert(
... external_id="my_agent",
... name="My agent",
... description="An agent with many tools",
... instructions="You are a helpful assistant that can query knowledge graphs, summarize documents, answer questions about documents, and query time series data points.",
... labels=["published"],
... tools=[find_assets_tool, find_files_tool, find_time_series_tool, summarize_tool, ask_doc_tool, ts_tool]
... tools=[
... find_assets_tool,
... find_files_tool,
... find_time_series_tool,
... summarize_tool,
... ask_doc_tool,
... ts_tool,
... ],
... )
>>> client.agents.upsert(agents=[agent])

Expand Down Expand Up @@ -281,8 +286,7 @@ async def chat(
>>> client = CogniteClient()
>>> # async_client = AsyncCogniteClient() # another option
>>> response = client.agents.chat(
... agent_external_id="my_agent",
... messages=Message("What can you help me with?")
... agent_external_id="my_agent", messages=Message("What can you help me with?")
... )
>>> print(response.text)

Expand All @@ -291,7 +295,7 @@ async def chat(
>>> follow_up = client.agents.chat(
... agent_external_id="my_agent",
... messages=Message("Tell me more about that"),
... cursor=response.cursor
... cursor=response.cursor,
... )

Send multiple messages at once:
Expand All @@ -300,8 +304,8 @@ async def chat(
... agent_external_id="my_agent",
... messages=[
... Message("Help me find the 1st stage compressor."),
... Message("Once you have found it, find related time series.")
... ]
... Message("Once you have found it, find related time series."),
... ],
... )

Chat with client-side actions:
Expand All @@ -316,13 +320,13 @@ async def chat(
... "a": {"type": "number", "description": "First number"},
... "b": {"type": "number", "description": "Second number"},
... },
... "required": ["a", "b"]
... }
... "required": ["a", "b"],
... },
... )
>>> response = client.agents.chat(
... agent_external_id="my_agent",
... messages=Message("What is 42 plus 58?"),
... actions=[add_numbers_action]
... actions=[add_numbers_action],
... )
>>> if response.action_calls:
... for call in response.action_calls:
Expand All @@ -332,11 +336,10 @@ async def chat(
... response = client.agents.chat(
... agent_external_id="my_agent",
... messages=ClientToolResult(
... action_id=call.action_id,
... content=f"The result is {result}"
... action_id=call.action_id, content=f"The result is {result}"
... ),
... cursor=response.cursor,
... actions=[add_numbers_action]
... actions=[add_numbers_action],
... )
"""
self._warnings.warn()
Expand Down
4 changes: 1 addition & 3 deletions cognite/client/_api/ai/tools/documents.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,7 @@ async def summarize(
You can also use external ID or instance ID:

>>> from cognite.client.data_classes.data_modeling import NodeId
>>> client.ai.tools.documents.summarize(
... instance_id=NodeId("my-space", "my-xid")
... )
>>> client.ai.tools.documents.summarize(instance_id=NodeId("my-space", "my-xid"))
"""
ident = IdentifierSequenceWithInstanceId.load(id, external_id, instance_id).as_singleton()
res = await self._post(
Expand Down
4 changes: 3 additions & 1 deletion cognite/client/_api/annotations.py
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,9 @@ async def list(self, filter: AnnotationFilter | dict, limit: int | None = DEFAUL
>>> from cognite.client.data_classes import AnnotationFilter
>>> client = CogniteClient()
>>> # async_client = AsyncCogniteClient() # another option
>>> flt = AnnotationFilter(annotated_resource_type="file", annotated_resource_ids=[{"id": 123}])
>>> flt = AnnotationFilter(
... annotated_resource_type="file", annotated_resource_ids=[{"id": 123}]
... )
>>> res = client.annotations.list(flt, limit=None)
"""
assert_type(filter, "filter", [AnnotationFilter, dict], allow_none=False)
Expand Down
65 changes: 45 additions & 20 deletions cognite/client/_api/assets.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,9 @@ async def retrieve_multiple(

Get assets by external id:

>>> res = client.assets.retrieve_multiple(external_ids=["abc", "def"], ignore_unknown_ids=True)
>>> res = client.assets.retrieve_multiple(
... external_ids=["abc", "def"], ignore_unknown_ids=True
... )
"""
identifiers = IdentifierSequence.load(ids=ids, external_ids=external_ids)
return await self._retrieve_multiple(
Expand Down Expand Up @@ -343,8 +345,8 @@ async def aggregate_cardinality_values(
>>> from cognite.client.data_classes.assets import AssetProperty
>>> is_critical = Search(AssetProperty.description, "critical")
>>> critical_assets = client.assets.aggregate_cardinality_values(
... AssetProperty.metadata_key("timezone"),
... advanced_filter=is_critical)
... AssetProperty.metadata_key("timezone"), advanced_filter=is_critical
... )
"""
self._validate_filter(advanced_filter)
return await self._advanced_aggregate(
Expand Down Expand Up @@ -421,7 +423,9 @@ async def aggregate_unique_values(
>>> from cognite.client.data_classes.assets import AssetProperty
>>> client = CogniteClient()
>>> # async_client = AsyncCogniteClient() # another option
>>> result = client.assets.aggregate_unique_values(AssetProperty.metadata_key("timezone"))
>>> result = client.assets.aggregate_unique_values(
... AssetProperty.metadata_key("timezone")
... )
>>> print(result.unique)

Get the different labels with count used for assets created after 2020-01-01 in your CDF project:
Expand All @@ -430,8 +434,12 @@ async def aggregate_unique_values(
>>> from cognite.client.data_classes.assets import AssetProperty
>>> from cognite.client.utils import timestamp_to_ms
>>> from datetime import datetime
>>> created_after_2020 = filters.Range(AssetProperty.created_time, gte=timestamp_to_ms(datetime(2020, 1, 1)))
>>> result = client.assets.aggregate_unique_values(AssetProperty.labels, advanced_filter=created_after_2020)
>>> created_after_2020 = filters.Range(
... AssetProperty.created_time, gte=timestamp_to_ms(datetime(2020, 1, 1))
... )
>>> result = client.assets.aggregate_unique_values(
... AssetProperty.labels, advanced_filter=created_after_2020
... )
>>> print(result.unique)

Get the different labels with count for assets updated after 2020-01-01 in your CDF project, but exclude all labels that
Expand All @@ -441,8 +449,14 @@ async def aggregate_unique_values(
>>> from cognite.client.data_classes import aggregations
>>> from cognite.client.data_classes import filters
>>> not_test = aggregations.Not(aggregations.Prefix("test"))
>>> created_after_2020 = filters.Range(AssetProperty.last_updated_time, gte=timestamp_to_ms(datetime(2020, 1, 1)))
>>> result = client.assets.aggregate_unique_values(AssetProperty.labels, advanced_filter=created_after_2020, aggregate_filter=not_test)
>>> created_after_2020 = filters.Range(
... AssetProperty.last_updated_time, gte=timestamp_to_ms(datetime(2020, 1, 1))
... )
>>> result = client.assets.aggregate_unique_values(
... AssetProperty.labels,
... advanced_filter=created_after_2020,
... aggregate_filter=not_test,
... )
>>> print(result.unique)

"""
Expand Down Expand Up @@ -600,7 +614,8 @@ async def create_hierarchy(
>>> assets = [
... AssetWrite(external_id="root", name="root"),
... AssetWrite(external_id="child1", parent_external_id="root", name="child1"),
... AssetWrite(external_id="child2", parent_external_id="root", name="child2")]
... AssetWrite(external_id="child2", parent_external_id="root", name="child2"),
... ]
>>> res = client.assets.create_hierarchy(assets)

Create an asset hierarchy, but run update for existing assets:
Expand Down Expand Up @@ -692,7 +707,7 @@ async def delete(
>>> from cognite.client import CogniteClient, AsyncCogniteClient
>>> client = CogniteClient()
>>> # async_client = AsyncCogniteClient() # another option
>>> client.assets.delete(id=[1,2,3], external_id="3")
>>> client.assets.delete(id=[1, 2, 3], external_id="3")
"""
await self._delete_multiple(
identifiers=IdentifierSequence.load(ids=id, external_ids=external_id),
Expand Down Expand Up @@ -735,7 +750,11 @@ async def update(
>>> from cognite.client.data_classes import AssetUpdate
>>> client = CogniteClient()
>>> # async_client = AsyncCogniteClient() # another option
>>> my_update = AssetUpdate(id=1).description.set("New description").metadata.add({"key": "value"})
>>> my_update = (
... AssetUpdate(id=1)
... .description.set("New description")
... .metadata.add({"key": "value"})
... )
>>> res1 = client.assets.update(my_update)
>>> # Remove an already set field like so
>>> another_update = AssetUpdate(id=1).description.set(None)
Expand Down Expand Up @@ -807,7 +826,9 @@ async def upsert(
>>> # async_client = AsyncCogniteClient() # another option
>>> existing_asset = client.assets.retrieve(id=1)
>>> existing_asset.description = "New description"
>>> new_asset = AssetWrite(external_id="new_asset", name="my asset", description="New asset")
>>> new_asset = AssetWrite(
... external_id="new_asset", name="my asset", description="New asset"
... )
>>> res = client.assets.upsert([existing_asset, new_asset], mode="replace")
"""
return await self._upsert_multiple(
Expand Down Expand Up @@ -862,12 +883,14 @@ async def search(

Search for assets using multiple filters, finding all assets with name similar to `xyz` with parent asset `123` or `456` with source `some source`:

>>> res = client.assets.search(name="xyz",filter={"parent_ids": [123,456],"source": "some source"})
>>> res = client.assets.search(
... name="xyz", filter={"parent_ids": [123, 456], "source": "some source"}
... )

Search for an asset with an attached label:

>>> my_label_filter = LabelFilter(contains_all=["PUMP"])
>>> res = client.assets.search(name="xyz",filter=AssetFilter(labels=my_label_filter))
>>> res = client.assets.search(name="xyz", filter=AssetFilter(labels=my_label_filter))
"""
return await self._search(
list_cls=AssetList,
Expand Down Expand Up @@ -988,7 +1011,7 @@ async def list(
Iterate over chunks of assets to reduce memory load:

>>> for asset_list in client.assets(chunk_size=2500):
... asset_list # do something with the assets
... asset_list # do something with the assets

Filter assets based on labels:

Expand All @@ -1013,17 +1036,19 @@ async def list(
>>> from cognite.client.data_classes.assets import AssetProperty, SortableAssetProperty
>>> in_timezone = filters.Prefix(AssetProperty.metadata_key("timezone"), "Europe")
>>> res = client.assets.list(
... advanced_filter=in_timezone,
... sort=(SortableAssetProperty.external_id, "asc"))
... advanced_filter=in_timezone, sort=(SortableAssetProperty.external_id, "asc")
... )

Combine filter and advanced filter:

>>> from cognite.client.data_classes import filters
>>> not_instrument_lvl5 = filters.And(
... filters.ContainsAny("labels", ["Level5"]),
... filters.Not(filters.ContainsAny("labels", ["Instrument"]))
... filters.ContainsAny("labels", ["Level5"]),
... filters.Not(filters.ContainsAny("labels", ["Instrument"])),
... )
>>> res = client.assets.list(
... asset_subtree_ids=[123456], advanced_filter=not_instrument_lvl5
... )
>>> res = client.assets.list(asset_subtree_ids=[123456], advanced_filter=not_instrument_lvl5)

"""
agg_props = self._process_aggregated_props(aggregated_properties)
Expand Down
Loading
Loading