feat: Google drive impersonator

.github/workflows/shared-packages.yml

@@ -200,6 +200,7 @@ jobs:
         env:
           GOOGLE_DRIVE_CLIENTID_JSON: ${{ secrets.GOOGLE_DRIVE_CLIENTID_JSON }}
           GOOGLE_SOURCE_UNIT_TEST_FOLDER: ${{ secrets.GOOGLE_SOURCE_UNIT_TEST_FOLDER }}
+          GOOGLE_DRIVE_TARGET_EMAIL: ${{ secrets.GOOGLE_DRIVE_TARGET_EMAIL }}
 
       - name: Test Report
         uses: mikepenz/action-junit-report@v4

docs/how-to/agents/define_and_use_agents.md

@@ -49,6 +49,14 @@ The result is an [AgentResult][ragbits.agents.AgentResult], which includes the m
 
 You can find the complete code example in the Ragbits repository [here](https://github.com/deepsense-ai/ragbits/blob/main/examples/agents/tool_use.py).
 
+## Tool choice
+To control what tool is used at first call you could use `tool_choice` parameter. There are the following options:
+- "auto": let model decide if tool call is needed
+- "none": do not call tool
+- "required: enforce tool usage (model decides which one)
+- Callable: one of provided tools
+
+
 ## Conversation history
 [`Agent`][ragbits.agents.Agent]s can retain conversation context across multiple interactions by enabling the `keep_history` flag when initializing the agent. This is useful when you want the agent to understand follow-up questions without needing the user to repeat earlier details.
 

docs/how-to/sources/google-drive.md

@@ -187,6 +187,43 @@ async def process_drive_documents():
 asyncio.run(process_drive_documents())
 ```
 
+## Impersonating Google Accounts
+
+You can configure your Google service account to impersonate other users in your Google Workspace domain. This is useful when you need to access files or perform actions on behalf of specific users.
+
+### Step 1: Enable Domain-Wide Delegation
+
+1. **Sign in to the [Google Admin Console](https://admin.google.com/) as a Super Admin.**
+2. Navigate to:
+    **Security > Access and data control > API controls > MANAGE DOMAIN WIDE DELEGATION**
+3. Add a new API client or edit an existing one, and include the following OAuth scopes:
+     - `https://www.googleapis.com/auth/cloud-platform`
+     - `https://www.googleapis.com/auth/drive`
+4. Click **Authorize** or **Save** to apply the changes.
+
+### Step 2: Impersonate a User in Your Code
+
+After configuring domain-wide delegation, you can specify a target user to impersonate when using the `GoogleDriveSource` in your code.
+
+```python
+from ragbits.core.sources.google_drive import GoogleDriveSource
+
+target_email = "[email protected]"
+credentials_file = "service-account-key.json"
+
+# Set the path to your service account key file
+GoogleDriveSource.set_credentials_file_path(credentials_file)
+
+# Set the email address of the user to impersonate
+GoogleDriveSource.set_impersonation_target(target_email)
+```
+
+**Note:**
+- The `target_email` must be a valid user in your Google Workspace domain.
+- Ensure your service account has been granted domain-wide delegation as described above.
+
+This setup allows your service account to act on behalf of the specified user, enabling access to their Google Drive files and resources as permitted by the assigned scopes.
+
 ## Troubleshooting
 
 ### Common Issues

examples/agents/tool_use.py

@@ -82,7 +82,7 @@ async def main() -> None:
         tools=[get_weather],
         default_options=AgentOptions(max_total_tokens=500, max_turns=5),
     )
-    response = await agent.run(WeatherPromptInput(location="Paris"))
+    response = await agent.run(WeatherPromptInput(location="Paris"), tool_choice=get_weather)
     print(response)
 
 

packages/ragbits-agents/CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
+- Add tool_choice parameter to agent interface (#738)
 
 ## 1.2.1 (2025-08-04)
 
@@ -13,7 +14,6 @@
 ### Changed
 
 - ragbits-core updated to version v1.2.0
-
 - Add native openai tools support (#621)
 - add Context to Agents (#715)
 

packages/ragbits-agents/src/ragbits/agents/_main.py

@@ -22,7 +22,7 @@
 )
 from ragbits.agents.mcp.server import MCPServer
 from ragbits.agents.mcp.utils import get_tools
-from ragbits.agents.tool import Tool, ToolCallResult
+from ragbits.agents.tool import Tool, ToolCallResult, ToolChoice
 from ragbits.core.audit.traces import trace
 from ragbits.core.llms.base import LLM, LLMClientOptionsT, LLMResponseWithMetadata, ToolCall, Usage
 from ragbits.core.options import Options
@@ -192,6 +192,7 @@ async def run(
         input: str | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResult[PromptOutputT]: ...
 
     @overload
@@ -200,13 +201,15 @@ async def run(
         input: PromptInputT,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResult[PromptOutputT]: ...
 
     async def run(
         self,
         input: str | PromptInputT | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResult[PromptOutputT]:
         """
         Run the agent. The method is experimental, inputs and outputs may change in the future.
@@ -218,6 +221,11 @@ async def run(
                 - None: No input. Only valid when a string prompt was provided during initialization.
             options: The options for the agent run.
             context: The context for the agent run.
+            tool_choice: Parameter that allows to control what tool is used at first call. Can be one of:
+                - "auto": let model decide if tool call is needed
+                - "none": do not call tool
+                - "required: enforce tool usage (model decides which one)
+                - Callable: one of provided tools
 
         Returns:
             The result of the agent run.
@@ -251,6 +259,7 @@ async def run(
                     await self.llm.generate_with_metadata(
                         prompt=prompt_with_history,
                         tools=[tool.to_function_schema() for tool in tools_mapping.values()],
+                        tool_choice=tool_choice if tool_choice and turn_count == 0 else None,
                         options=self._get_llm_options(llm_options, merged_options, context.usage),
                     ),
                 )
@@ -294,6 +303,7 @@ def run_streaming(
         input: str | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResultStreaming: ...
 
     @overload
@@ -302,13 +312,15 @@ def run_streaming(
         input: PromptInputT,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResultStreaming: ...
 
     def run_streaming(
         self,
         input: str | PromptInputT | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResultStreaming:
         """
         This method returns an `AgentResultStreaming` object that can be asynchronously
@@ -318,6 +330,11 @@ def run_streaming(
             input: The input for the agent run.
             options: The options for the agent run.
             context: The context for the agent run.
+            tool_choice: Parameter that allows to control what tool is used at first call. Can be one of:
+                - "auto": let model decide if tool call is needed
+                - "none": do not call tool
+                - "required: enforce tool usage (model decides which one)
+                - Callable: one of provided tools
 
         Returns:
             A `StreamingResult` object for iteration and collection.
@@ -329,14 +346,15 @@ def run_streaming(
             AgentInvalidPromptInputError: If the prompt/input combination is invalid.
             AgentMaxTurnsExceededError: If the maximum number of turns is exceeded.
         """
-        generator = self._stream_internal(input, options, context)
+        generator = self._stream_internal(input, options, context, tool_choice)
         return AgentResultStreaming(generator)
 
     async def _stream_internal(
         self,
         input: str | PromptInputT | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AsyncGenerator[str | ToolCall | ToolCallResult | SimpleNamespace | BasePrompt | Usage]:
         if context is None:
             context = AgentRunContext()
@@ -357,6 +375,7 @@ async def _stream_internal(
                 streaming_result = self.llm.generate_streaming(
                     prompt=prompt_with_history,
                     tools=[tool.to_function_schema() for tool in tools_mapping.values()],
+                    tool_choice=tool_choice if tool_choice and turn_count == 0 else None,
                     options=self._get_llm_options(llm_options, merged_options, context.usage),
                 )
                 async for chunk in streaming_result:

packages/ragbits-agents/src/ragbits/agents/tool.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 from typing_extensions import Self
 
@@ -76,3 +76,6 @@ def to_function_schema(self) -> dict[str, Any]:
                 "parameters": self.parameters,
             },
         }
+
+
+ToolChoice = Literal["auto", "none", "required"] | Callable