From 22cfdd6ae3bb396469d48dcc93564efae0254e71 Mon Sep 17 00:00:00 2001
From: rafaelmmiller <150964962+rafaelsideguide@users.noreply.github.com>
Date: Thu, 17 Apr 2025 23:31:28 -0700
Subject: [PATCH] added agent options types

---
 apps/python-sdk/firecrawl/firecrawl.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/apps/python-sdk/firecrawl/firecrawl.py b/apps/python-sdk/firecrawl/firecrawl.py
index 8ca94fa6..d7abf336 100644
--- a/apps/python-sdk/firecrawl/firecrawl.py
+++ b/apps/python-sdk/firecrawl/firecrawl.py
@@ -76,6 +76,15 @@ class FirecrawlDocumentMetadata(pydantic.BaseModel):
     statusCode: Optional[int] = None
     error: Optional[str] = None
 
+class AgentOptions(pydantic.BaseModel):
+    """Configuration for the agent."""
+    model: Literal["FIRE-1"] = "FIRE-1"
+    prompt: Optional[str] = None
+
+class AgentOptionsExtract(pydantic.BaseModel):
+    """Configuration for the agent in extract operations."""
+    model: Literal["FIRE-1"] = "FIRE-1"
+
 class ActionsResult(pydantic.BaseModel):
     """Result of actions performed during scraping."""
     screenshots: List[str]
@@ -164,17 +173,24 @@ class ExecuteJavascriptAction(pydantic.BaseModel):
     type: Literal["executeJavascript"]
     script: str
 
+
+class ExtractAgent(pydantic.BaseModel):
+    """Configuration for the agent in extract operations."""
+    model: Literal["FIRE-1"] = "FIRE-1"
+
 class ExtractConfig(pydantic.BaseModel):
     """Configuration for extraction."""
     prompt: Optional[str] = None
     schema: Optional[Any] = None
     systemPrompt: Optional[str] = None
+    agent: Optional[ExtractAgent] = None
 
 class ScrapeParams(CrawlScrapeOptions):
     """Parameters for scraping operations."""
     extract: Optional[ExtractConfig] = None
     jsonOptions: Optional[ExtractConfig] = None
     actions: Optional[List[Union[WaitAction, ScreenshotAction, ClickAction, WriteAction, PressAction, ScrollAction, ScrapeAction, ExecuteJavascriptAction]]] = None
+    agent: Optional[AgentOptions] = None
 
 class ScrapeResponse(FirecrawlDocument[T], Generic[T]):
     """Response from scraping operations."""
@@ -363,7 +379,7 @@ class ChangeTrackingData(pydantic.BaseModel):
     visibility: str  # "visible" | "hidden"
     diff: Optional[Dict[str, Any]] = None
     json: Optional[Any] = None
-
+    
 class FirecrawlApp:
     class SearchResponse(pydantic.BaseModel):
         """