Excavator: Upgrade API Version (#106)

palantir · Jan 23, 2025 · 4d6439f · 4d6439f
1 parent 2fce34a
commit 4d6439f
Show file tree

Hide file tree

Showing 22 changed files with 148 additions and 31 deletions.
diff --git a/README.md b/README.md
@@ -659,6 +659,8 @@ Namespace | Resource | Operation | HTTP request |
 - [CancelSessionResponseDict](docs/v2/models/CancelSessionResponseDict.md)
 - [Content](docs/v2/models/Content.md)
 - [ContentDict](docs/v2/models/ContentDict.md)
+- [FunctionRetrievedContext](docs/v2/models/FunctionRetrievedContext.md)
+- [FunctionRetrievedContextDict](docs/v2/models/FunctionRetrievedContextDict.md)
 - [InputContextDict](docs/v2/models/InputContextDict.md)
 - [ListAgentVersionsResponse](docs/v2/models/ListAgentVersionsResponse.md)
 - [ListAgentVersionsResponseDict](docs/v2/models/ListAgentVersionsResponseDict.md)

diff --git a/docs/v2/Datasets/Branch.md b/docs/v2/Datasets/Branch.md
@@ -18,7 +18,7 @@ Name | Type | Description  | Notes |
 ------------- | ------------- | ------------- | ------------- |
 **dataset_rid** | DatasetRid | datasetRid |  |
 **name** | BranchName |  |  |
-**transaction_rid** | Optional[TransactionRid] |  | [optional] |
+**transaction_rid** | Optional[TransactionRid] | The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction. | [optional] |
 
 ### Return type
 **Branch**
@@ -38,7 +38,7 @@ foundry_client = FoundryClient(
 dataset_rid = None
 # BranchName |
 name = "master"
-# Optional[TransactionRid] |
+# Optional[TransactionRid] | The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction.
 transaction_rid = "ri.foundry.main.transaction.0a0207cb-26b7-415b-bc80-66a3aa3933f4"
 
 

diff --git a/docs/v2/aip_agents/models/AgentSessionRagContextResponse.md b/docs/v2/aip_agents/models/AgentSessionRagContextResponse.md
@@ -7,6 +7,7 @@ Context retrieved from an Agent's configured context data sources which was rele
 | Name | Type | Required | Description |
 | ------------ | ------------- | ------------- | ------------- |
 **object_contexts** | List[ObjectContext] | Yes |  |
+**function_retrieved_contexts** | List[FunctionRetrievedContext] | Yes |  |
 
 
 [[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/docs/v2/aip_agents/models/AgentSessionRagContextResponseDict.md b/docs/v2/aip_agents/models/AgentSessionRagContextResponseDict.md
@@ -7,6 +7,7 @@ Context retrieved from an Agent's configured context data sources which was rele
 | Name | Type | Required | Description |
 | ------------ | ------------- | ------------- | ------------- |
 **objectContexts** | List[ObjectContextDict] | Yes |  |
+**functionRetrievedContexts** | List[FunctionRetrievedContextDict] | Yes |  |
 
 
 [[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/docs/v2/aip_agents/models/FunctionRetrievedContext.md b/docs/v2/aip_agents/models/FunctionRetrievedContext.md
@@ -0,0 +1,15 @@
+# FunctionRetrievedContext
+
+Context retrieved from running a function to include as additional context in the prompt to the Agent.
+
+
+## Properties
+| Name | Type | Required | Description |
+| ------------ | ------------- | ------------- | ------------- |
+**function_rid** | FunctionRid | Yes |  |
+**function_version** | FunctionVersion | Yes |  |
+**retrieved_prompt** | str | Yes | String content returned from a context retrieval function.  |
+**type** | Literal["functionRetrievedContext"] | Yes | None |
+
+
+[[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/docs/v2/aip_agents/models/FunctionRetrievedContextDict.md b/docs/v2/aip_agents/models/FunctionRetrievedContextDict.md
@@ -0,0 +1,15 @@
+# FunctionRetrievedContextDict
+
+Context retrieved from running a function to include as additional context in the prompt to the Agent.
+
+
+## Properties
+| Name | Type | Required | Description |
+| ------------ | ------------- | ------------- | ------------- |
+**functionRid** | FunctionRid | Yes |  |
+**functionVersion** | FunctionVersion | Yes |  |
+**retrievedPrompt** | str | Yes | String content returned from a context retrieval function.  |
+**type** | Literal["functionRetrievedContext"] | Yes | None |
+
+
+[[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/docs/v2/aip_agents/models/InputContextDict.md b/docs/v2/aip_agents/models/InputContextDict.md
@@ -3,10 +3,15 @@
 Custom retrieved [context](/docs/foundry/agent-studio/retrieval-context/) to provide to an Agent for continuing a session.
 
 
-## Type
-```python
-ObjectContextDict
-```
+This is a discriminator type and does not contain any fields. Instead, it is a union
+of of the models listed below.
+
+This discriminator class uses the `type` field to differentiate between classes.
+
+| Class | Value
+| ------------ | -------------
+FunctionRetrievedContextDict | functionRetrievedContext
+ObjectContextDict | objectContext
 
 
 [[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/docs/v2/datasets/models/Branch.md b/docs/v2/datasets/models/Branch.md
@@ -6,7 +6,7 @@ Branch
 | Name | Type | Required | Description |
 | ------------ | ------------- | ------------- | ------------- |
 **name** | BranchName | Yes |  |
-**transaction_rid** | Optional[TransactionRid] | No |  |
+**transaction_rid** | Optional[TransactionRid] | No | The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction. |
 
 
 [[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/docs/v2/datasets/models/BranchDict.md b/docs/v2/datasets/models/BranchDict.md
@@ -6,7 +6,7 @@ Branch
 | Name | Type | Required | Description |
 | ------------ | ------------- | ------------- | ------------- |
 **name** | BranchName | Yes |  |
-**transactionRid** | NotRequired[TransactionRid] | No |  |
+**transactionRid** | NotRequired[TransactionRid] | No | The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction. |
 
 
 [[Back to Model list]](../../../../README.md#models-v2-link) [[Back to API list]](../../../../README.md#apis-v2-link) [[Back to README]](../../../../README.md)
diff --git a/foundry/_versions.py b/foundry/_versions.py
@@ -17,4 +17,4 @@
 # using the autorelease bot
 __version__ = "0.0.0"
 
-__openapi_document_version__ = "1.1043.0"
+__openapi_document_version__ = "1.1044.0"
diff --git a/foundry/v2/aip_agents/models/__init__.py b/foundry/v2/aip_agents/models/__init__.py
@@ -38,6 +38,12 @@
 )  # NOQA
 from foundry.v2.aip_agents.models._content import Content
 from foundry.v2.aip_agents.models._content_dict import ContentDict
+from foundry.v2.aip_agents.models._function_retrieved_context import (
+    FunctionRetrievedContext,
+)  # NOQA
+from foundry.v2.aip_agents.models._function_retrieved_context_dict import (
+    FunctionRetrievedContextDict,
+)  # NOQA
 from foundry.v2.aip_agents.models._input_context_dict import InputContextDict
 from foundry.v2.aip_agents.models._list_agent_versions_response import (
     ListAgentVersionsResponse,
@@ -118,6 +124,8 @@
     "CancelSessionResponseDict",
     "Content",
     "ContentDict",
+    "FunctionRetrievedContext",
+    "FunctionRetrievedContextDict",
     "InputContextDict",
     "ListAgentVersionsResponse",
     "ListAgentVersionsResponseDict",

diff --git a/foundry/v2/aip_agents/models/_agent_session_rag_context_response.py b/foundry/v2/aip_agents/models/_agent_session_rag_context_response.py
@@ -23,6 +23,9 @@
 from foundry.v2.aip_agents.models._agent_session_rag_context_response_dict import (
     AgentSessionRagContextResponseDict,
 )  # NOQA
+from foundry.v2.aip_agents.models._function_retrieved_context import (
+    FunctionRetrievedContext,
+)  # NOQA
 from foundry.v2.aip_agents.models._object_context import ObjectContext
 
 
@@ -31,6 +34,10 @@ class AgentSessionRagContextResponse(pydantic.BaseModel):
 
     object_contexts: List[ObjectContext] = pydantic.Field(alias="objectContexts")
 
+    function_retrieved_contexts: List[FunctionRetrievedContext] = pydantic.Field(
+        alias="functionRetrievedContexts"
+    )
+
     model_config = {"extra": "allow"}
 
     def to_dict(self) -> AgentSessionRagContextResponseDict:

diff --git a/foundry/v2/aip_agents/models/_agent_session_rag_context_response_dict.py b/foundry/v2/aip_agents/models/_agent_session_rag_context_response_dict.py
@@ -19,6 +19,9 @@
 
 from typing_extensions import TypedDict
 
+from foundry.v2.aip_agents.models._function_retrieved_context_dict import (
+    FunctionRetrievedContextDict,
+)  # NOQA
 from foundry.v2.aip_agents.models._object_context_dict import ObjectContextDict
 
 
@@ -28,3 +31,5 @@ class AgentSessionRagContextResponseDict(TypedDict):
     __pydantic_config__ = {"extra": "allow"}  # type: ignore
 
     objectContexts: List[ObjectContextDict]
+
+    functionRetrievedContexts: List[FunctionRetrievedContextDict]
diff --git a/foundry/v2/aip_agents/models/_function_retrieved_context.py b/foundry/v2/aip_agents/models/_function_retrieved_context.py
@@ -0,0 +1,49 @@
+#  Copyright 2024 Palantir Technologies, Inc.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+
+from __future__ import annotations
+
+from typing import Literal
+from typing import cast
+
+import pydantic
+
+from foundry.v2.aip_agents.models._function_retrieved_context_dict import (
+    FunctionRetrievedContextDict,
+)  # NOQA
+from foundry.v2.functions.models._function_rid import FunctionRid
+from foundry.v2.functions.models._function_version import FunctionVersion
+
+
+class FunctionRetrievedContext(pydantic.BaseModel):
+    """Context retrieved from running a function to include as additional context in the prompt to the Agent."""
+
+    function_rid: FunctionRid = pydantic.Field(alias="functionRid")
+
+    function_version: FunctionVersion = pydantic.Field(alias="functionVersion")
+
+    retrieved_prompt: str = pydantic.Field(alias="retrievedPrompt")
+
+    """String content returned from a context retrieval function."""
+
+    type: Literal["functionRetrievedContext"] = "functionRetrievedContext"
+
+    model_config = {"extra": "allow"}
+
+    def to_dict(self) -> FunctionRetrievedContextDict:
+        """Return the dictionary representation of the model using the field aliases."""
+        return cast(
+            FunctionRetrievedContextDict, self.model_dump(by_alias=True, exclude_unset=True)
+        )
diff --git a/...s/errors/_read_table_permission_denied.py → ...odels/_function_retrieved_context_dict.py b/...s/errors/_read_table_permission_denied.py → ...odels/_function_retrieved_context_dict.py
@@ -15,29 +15,24 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
 from typing import Literal
 
-import pydantic
 from typing_extensions import TypedDict
 
-from foundry._errors import PalantirRPCException
-from foundry.v2.datasets.models._dataset_rid import DatasetRid
+from foundry.v2.functions.models._function_rid import FunctionRid
+from foundry.v2.functions.models._function_version import FunctionVersion
 
 
-class ReadTablePermissionDeniedParameters(TypedDict):
-    """The provided token does not have permission to read the given dataset as a table."""
+class FunctionRetrievedContextDict(TypedDict):
+    """Context retrieved from running a function to include as additional context in the prompt to the Agent."""
 
     __pydantic_config__ = {"extra": "allow"}  # type: ignore
 
-    datasetRid: DatasetRid
+    functionRid: FunctionRid
 
+    functionVersion: FunctionVersion
 
-@dataclass
-class ReadTablePermissionDenied(PalantirRPCException):
-    name: Literal["ReadTablePermissionDenied"]
-    parameters: ReadTablePermissionDeniedParameters
-    error_instance_id: str
+    retrievedPrompt: str
+    """String content returned from a context retrieval function."""
 
-
-__all__ = ["ReadTablePermissionDenied"]
+    type: Literal["functionRetrievedContext"]
diff --git a/foundry/v2/aip_agents/models/_input_context_dict.py b/foundry/v2/aip_agents/models/_input_context_dict.py
@@ -15,7 +15,17 @@
 
 from __future__ import annotations
 
+from typing import Union
+
+import pydantic
+from typing_extensions import Annotated
+
+from foundry.v2.aip_agents.models._function_retrieved_context_dict import (
+    FunctionRetrievedContextDict,
+)  # NOQA
 from foundry.v2.aip_agents.models._object_context_dict import ObjectContextDict
 
-InputContextDict = ObjectContextDict
+InputContextDict = Annotated[
+    Union[FunctionRetrievedContextDict, ObjectContextDict], pydantic.Field(discriminator="type")
+]
 """Custom retrieved [context](/docs/foundry/agent-studio/retrieval-context/) to provide to an Agent for continuing a session."""
diff --git a/foundry/v2/cli.py b/foundry/v2/cli.py
@@ -2510,7 +2510,12 @@ def datasets_dataset_branch():
 @datasets_dataset_branch.command("create")
 @click.argument("dataset_rid", type=str, required=True)
 @click.option("--name", type=str, required=True, help="""""")
-@click.option("--transaction_rid", type=str, required=False, help="""""")
+@click.option(
+    "--transaction_rid",
+    type=str,
+    required=False,
+    help="""The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction.""",
+)
 @click.pass_obj
 def datasets_dataset_branch_create(
     client: foundry.v2.FoundryClient,

diff --git a/foundry/v2/datasets/branch.py b/foundry/v2/datasets/branch.py
@@ -75,7 +75,7 @@ def create(
         :type dataset_rid: DatasetRid
         :param name:
         :type name: BranchName
-        :param transaction_rid:
+        :param transaction_rid: The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction.
         :type transaction_rid: Optional[TransactionRid]
         :param request_timeout: timeout setting for this request in seconds.
         :type request_timeout: Optional[int]

diff --git a/foundry/v2/datasets/errors/__init__.py b/foundry/v2/datasets/errors/__init__.py
@@ -62,9 +62,6 @@
 from foundry.v2.datasets.errors._read_table_dataset_permission_denied import (
     ReadTableDatasetPermissionDenied,
 )  # NOQA
-from foundry.v2.datasets.errors._read_table_permission_denied import (
-    ReadTablePermissionDenied,
-)  # NOQA
 from foundry.v2.datasets.errors._schema_not_found import SchemaNotFound
 from foundry.v2.datasets.errors._transaction_not_committed import TransactionNotCommitted  # NOQA
 from foundry.v2.datasets.errors._transaction_not_found import TransactionNotFound
@@ -97,7 +94,6 @@
     "OpenTransactionAlreadyExists",
     "PutSchemaPermissionDenied",
     "ReadTableDatasetPermissionDenied",
-    "ReadTablePermissionDenied",
     "SchemaNotFound",
     "TransactionNotCommitted",
     "TransactionNotFound",

diff --git a/foundry/v2/datasets/errors/_read_table_dataset_permission_denied.py b/foundry/v2/datasets/errors/_read_table_dataset_permission_denied.py
@@ -26,7 +26,7 @@
 
 
 class ReadTableDatasetPermissionDeniedParameters(TypedDict):
-    """Could not readTable the Dataset."""
+    """The provided token does not have permission to read the given dataset as a table."""
 
     __pydantic_config__ = {"extra": "allow"}  # type: ignore
 

diff --git a/foundry/v2/datasets/models/_branch.py b/foundry/v2/datasets/models/_branch.py
@@ -32,6 +32,8 @@ class Branch(pydantic.BaseModel):
 
     transaction_rid: Optional[TransactionRid] = pydantic.Field(alias="transactionRid", default=None)
 
+    """The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction."""
+
     model_config = {"extra": "allow"}
 
     def to_dict(self) -> BranchDict:

diff --git a/foundry/v2/datasets/models/_branch_dict.py b/foundry/v2/datasets/models/_branch_dict.py
@@ -30,3 +30,4 @@ class BranchDict(TypedDict):
     name: BranchName
 
     transactionRid: NotRequired[TransactionRid]
+    """The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction."""
Original file line number	Diff line number	Diff line change
Expand Up		@@ -30,3 +30,4 @@ class BranchDict(TypedDict):
		name: BranchName

		transactionRid: NotRequired[TransactionRid]
		"""The most recent OPEN or COMMITTED transaction on the branch. This will never be an ABORTED transaction."""