From b08f7a9765e25449ba7e7695bc2698ba6eac350c Mon Sep 17 00:00:00 2001 From: aryan Date: Mon, 14 Apr 2025 00:58:05 +0530 Subject: [PATCH 01/24] feat(model-preference): Implement model preference selection across the application - Added model preference handling in SystemInstructor and WebSurfer classes. - Introduced API endpoints for setting and retrieving model preference. - Updated FastAPI app to manage model preference during websocket connections and response generation. - Enhanced orchestrator and session management to utilize model preference. - Integrated model preference into frontend components for user selection. --- cortex_on/agents/orchestrator_agent.py | 8 +- cortex_on/agents/web_surfer.py | 10 +- cortex_on/instructor.py | 6 +- cortex_on/main.py | 62 +++++++++-- frontend/src/components/home/Header.tsx | 12 ++- frontend/src/components/ui/ModelToggle.tsx | 61 +++++++++++ .../src/dataStore/modelPreferenceSlice.ts | 22 ++++ frontend/src/dataStore/store.ts | 7 +- frontend/src/lib/utils.ts | 4 +- frontend/src/services/modelPreferenceApi.ts | 23 ++++ .../task_20250413_190916/conversation.json | 79 ++++++++++++++ .../task_20250413_191154/conversation.json | 1 + .../task_20250413_192014/conversation.json | 49 +++++++++ ta-browser/core/orchestrator.py | 52 ++++----- ta-browser/core/server/models/web.py | 1 + ta-browser/core/server/routes/web.py | 7 +- .../core/server/utils/session_tracker.py | 11 +- ta-browser/core/utils/init_client.py | 102 +++++++++++++----- ta-browser/core/utils/openai_client.py | 65 +---------- 19 files changed, 441 insertions(+), 141 deletions(-) create mode 100644 frontend/src/components/ui/ModelToggle.tsx create mode 100644 frontend/src/dataStore/modelPreferenceSlice.ts create mode 100644 frontend/src/services/modelPreferenceApi.ts create mode 100644 ta-browser/1744571356739037821/task_20250413_190916/conversation.json create mode 100644 ta-browser/1744571514482277604/task_20250413_191154/conversation.json create mode 100644 ta-browser/1744572014568045555/task_20250413_192014/conversation.json diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 6b001ba..284d927 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -22,6 +22,7 @@ class orchestrator_deps: stream_output: Optional[StreamResponse] = None # Add a collection to track agent-specific streams agent_responses: Optional[List[StreamResponse]] = None + model_preference: str = "Anthropic" orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. @@ -307,8 +308,11 @@ async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - # Initialize WebSurfer agent - web_surfer_agent = WebSurfer(api_url="http://localhost:8000/api/v1/web/stream") + # Initialize WebSurfer agent with model preference + web_surfer_agent = WebSurfer( + api_url="http://localhost:8000/api/v1/web/stream", + model_preference=ctx.deps.model_preference + ) # Run WebSurfer with its own stream_output success, message, messages = await web_surfer_agent.generate_reply( diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 34e2cfd..03503d5 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -29,19 +29,25 @@ TIMEOUT = 9999999999999999999999999999999999999999999 class WebSurfer: - def __init__(self, api_url: str = "http://localhost:8000/api/v1/web/stream"): + def __init__(self, api_url: str = "http://localhost:8000/api/v1/web/stream", model_preference: str = "Anthropic"): self.api_url = api_url self.name = "Web Surfer Agent" self.description = "An agent that is a websurfer and a webscraper that can access any web-page to extract information or perform actions." self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None + self.model_preference = model_preference async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, Any]]]: session_timeout = aiohttp.ClientTimeout(total=None, sock_connect=TIMEOUT, sock_read=TIMEOUT) async with aiohttp.ClientSession(timeout=session_timeout) as session: final_json_response = [] try: - payload = {"cmd": instruction, "critique_disabled": False} + payload = { + "cmd": instruction, + "critique_disabled": False, + "model_preference": self.model_preference + } + print(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index b4f0efb..f2688b4 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -38,10 +38,11 @@ def default(self, obj): # Main Orchestrator Class class SystemInstructor: - def __init__(self): + def __init__(self, model_preference: str = "Anthropic"): self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None self.orchestrator_response: List[StreamResponse] = [] + self.model_preference = model_preference self._setup_logging() def _setup_logging(self) -> None: @@ -80,7 +81,8 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: deps_for_orchestrator = orchestrator_deps( websocket=self.websocket, stream_output=stream_output, - agent_responses=self.orchestrator_response # Pass reference to collection + agent_responses=self.orchestrator_response, # Pass reference to collection + model_preference=self.model_preference ) try: diff --git a/cortex_on/main.py b/cortex_on/main.py index a8dd4de..e1adddd 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -1,27 +1,75 @@ # Standard library imports from typing import List, Optional +from contextlib import asynccontextmanager # Third-party imports -from fastapi import FastAPI, WebSocket +from fastapi import FastAPI, WebSocket, Depends +from fastapi.middleware.cors import CORSMiddleware # Local application imports from instructor import SystemInstructor -app: FastAPI = FastAPI() +# Default model preference is Anthropic +MODEL_PREFERENCE = "Anthropic" -async def generate_response(task: str, websocket: Optional[WebSocket] = None): - orchestrator: SystemInstructor = SystemInstructor() +@asynccontextmanager +async def lifespan(app: FastAPI): + # Set default model preference at startup + app.state.model_preference = MODEL_PREFERENCE + print(f"[STARTUP] Setting default model preference to: {MODEL_PREFERENCE}") + yield + + +app: FastAPI = FastAPI(lifespan=lifespan) + +# Add CORS middleware +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], # Allow all origins for development + allow_credentials=True, + allow_methods=["*"], # Allow all methods + allow_headers=["*"], # Allow all headers +) + +async def get_model_preference() -> str: + """ + Get the current model preference from app state + """ + print(f"[GET] Current model preference: {app.state.model_preference}") + return app.state.model_preference + +@app.get("/set_model_preference") +async def set_model_preference(model: str): + """ + Set the model preference (Anthropic or OpenAI) + """ + if model not in ["Anthropic", "OpenAI"]: + print(f"[ERROR] Invalid model preference attempted: {model}") + return {"error": "Model must be 'Anthropic' or 'OpenAI'"} + print(f"[SET] Changing model preference from {app.state.model_preference} to {model}") + app.state.model_preference = model + return {"message": f"Model preference set to {model}"} + +async def generate_response(task: str, websocket: Optional[WebSocket] = None, model_preference: str = None): + if model_preference is None: + model_preference = app.state.model_preference + print(f"[GENERATE] Using model preference: {model_preference} for task: {task[:30]}...") + orchestrator: SystemInstructor = SystemInstructor(model_preference) return await orchestrator.run(task, websocket) @app.get("/agent/chat") -async def agent_chat(task: str) -> List: - final_agent_response = await generate_response(task) +async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List: + print(f"[CHAT] Received chat request with model preference: {model_preference}") + final_agent_response = await generate_response(task, model_preference=model_preference) return final_agent_response @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() + model_preference = app.state.model_preference + print(f"[WEBSOCKET] New connection using model preference: {model_preference}") while True: data = await websocket.receive_text() - await generate_response(data, websocket) + print(f"[WEBSOCKET] Received message, using model: {model_preference}") + await generate_response(data, websocket, model_preference) diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index d6217fd..5fba65a 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -4,6 +4,7 @@ import {useDispatch} from "react-redux"; import {useLocation, useNavigate} from "react-router-dom"; import Logo from "../../assets/CortexON_logo_dark.svg"; import {Button} from "../ui/button"; +import ModelToggle from "../ui/ModelToggle"; const Header = () => { const nav = useNavigate(); @@ -13,7 +14,7 @@ const Header = () => { return (
{ dispatch(setMessages([])); nav("/"); @@ -21,10 +22,11 @@ const Header = () => { > Logo
-
+ +
nav("/vault")} - className={`w-[10%] h-full flex justify-center items-center cursor-pointer border-b-2 hover:border-[#BD24CA] ${ + className={`h-full flex justify-center items-center cursor-pointer border-b-2 px-4 hover:border-[#BD24CA] ${ location.includes("/vault") ? "border-[#BD24CA]" : "border-background" @@ -32,7 +34,11 @@ const Header = () => { >

Vault

+ + {/* Model Toggle Button */} +
+ + + + GPT-4o + +
+ ); +} \ No newline at end of file diff --git a/frontend/src/dataStore/modelPreferenceSlice.ts b/frontend/src/dataStore/modelPreferenceSlice.ts new file mode 100644 index 0000000..7c1d602 --- /dev/null +++ b/frontend/src/dataStore/modelPreferenceSlice.ts @@ -0,0 +1,22 @@ +import { createSlice, PayloadAction } from "@reduxjs/toolkit"; + +interface ModelPreferenceState { + preference: "Anthropic" | "OpenAI"; +} + +const initialState: ModelPreferenceState = { + preference: "Anthropic", +}; + +const modelPreferenceSlice = createSlice({ + name: "modelPreference", + initialState, + reducers: { + setModelPreference: (state, action: PayloadAction<"Anthropic" | "OpenAI">) => { + state.preference = action.payload; + }, + }, +}); + +export const { setModelPreference } = modelPreferenceSlice.actions; +export default modelPreferenceSlice; \ No newline at end of file diff --git a/frontend/src/dataStore/store.ts b/frontend/src/dataStore/store.ts index c3be4e6..516535f 100644 --- a/frontend/src/dataStore/store.ts +++ b/frontend/src/dataStore/store.ts @@ -1,13 +1,18 @@ import vaultApi from "@/services/vaultApi"; +import modelPreferenceApi from "@/services/modelPreferenceApi"; import {configureStore} from "@reduxjs/toolkit"; import messagesSlice from "./messagesSlice"; +import modelPreferenceSlice from "./modelPreferenceSlice"; + export const store = configureStore({ reducer: { [messagesSlice.name]: messagesSlice.reducer, [vaultApi.reducerPath]: vaultApi.reducer, + [modelPreferenceSlice.name]: modelPreferenceSlice.reducer, + [modelPreferenceApi.reducerPath]: modelPreferenceApi.reducer, }, middleware: (getDefaultMiddleware) => - getDefaultMiddleware().concat(vaultApi.middleware), + getDefaultMiddleware().concat(vaultApi.middleware, modelPreferenceApi.middleware), }); export type RootState = ReturnType; diff --git a/frontend/src/lib/utils.ts b/frontend/src/lib/utils.ts index dfe0774..07ab4f7 100644 --- a/frontend/src/lib/utils.ts +++ b/frontend/src/lib/utils.ts @@ -1,5 +1,5 @@ -import {clsx, type ClassValue} from "clsx"; -import {twMerge} from "tailwind-merge"; +import { type ClassValue, clsx } from "clsx"; +import { twMerge } from "tailwind-merge"; export function cn(...inputs: ClassValue[]) { return twMerge(clsx(inputs)); diff --git a/frontend/src/services/modelPreferenceApi.ts b/frontend/src/services/modelPreferenceApi.ts new file mode 100644 index 0000000..e7320a6 --- /dev/null +++ b/frontend/src/services/modelPreferenceApi.ts @@ -0,0 +1,23 @@ +import { createApi } from "@reduxjs/toolkit/query/react"; +import { fetchBaseQuery } from "@reduxjs/toolkit/query/react"; + +// Get the base URL from environment variable or use a default +const CORTEX_ON_API_URL = import.meta.env.VITE_CORTEX_ON_API_URL || "http://localhost:8081"; + +const modelPreferenceApi = createApi({ + reducerPath: "modelPreferenceApi", + baseQuery: fetchBaseQuery({ + baseUrl: CORTEX_ON_API_URL, + }), + endpoints: (builder) => ({ + setModelPreference: builder.mutation<{ message: string }, string>({ + query: (model) => ({ + url: `/set_model_preference?model=${model}`, + method: "GET", + }), + }), + }), +}); + +export const { useSetModelPreferenceMutation } = modelPreferenceApi; +export default modelPreferenceApi; \ No newline at end of file diff --git a/ta-browser/1744571356739037821/task_20250413_190916/conversation.json b/ta-browser/1744571356739037821/task_20250413_190916/conversation.json new file mode 100644 index 0000000..e9976ca --- /dev/null +++ b/ta-browser/1744571356739037821/task_20250413_190916/conversation.json @@ -0,0 +1,79 @@ +[ + { + "role": "user", + "content": "Inputs :\nUser Query : Go to Amazon.com and search for RTX 4060 graphics cards. Return the search results including prices and available models.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" + }, + { + "role": "assistant", + "content": "Plan: 1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\nNext Step: Navigate to Amazon's website (https://www.amazon.com)", + "name": "planner_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" current_step=\"Navigate to Amazon's website (https://www.amazon.com)\" " + }, + { + "role": "assistant", + "content": "I'll help navigate to Amazon's website using the open_url_tool.", + "name": "browser_nav_agent" + }, + { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", + "type": "function", + "function": { + "name": "open_url_tool", + "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", + "name": "open_url_tool", + "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" + }, + { + "role": "assistant", + "content": "I attempted to navigate to Amazon's website but encountered an error. This is common with Amazon as they have strict bot detection measures. Let me verify if we landed on any page by checking the current URL.", + "name": "browser_nav_agent" + }, + { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", + "type": "function", + "function": { + "name": "get_url", + "arguments": "{\"raw_args\": \"ArgsDict(args_dict={})\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", + "name": "get_url", + "content": "Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more." + }, + { + "role": "assistant", + "content": "Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.", + "name": "browser_nav_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" next_step=\"Navigate to Amazon's website (https://www.amazon.com)\" tool_response=\"Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\nTool Call: get_url\nArguments: {}\nResponse: Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more.\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" + }, + { + "role": "assistant", + "content": "Feedback: Original Plan:\n1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\n\nCurrent Progress:\nCurrently on Step 1: Navigate to Amazon's website (https://www.amazon.com)\n\nAnalysis and Feedback:\nThe browser agent has successfully completed the first step of navigating to Amazon.com. This is confirmed by:\n1. The get_url tool response showing we are at \"https://www.amazon.com/\"\n2. The page title verification showing \"Amazon.com. Spend less. Smile more.\"\n\nDespite an initial execution context error (which is common during page loads), the navigation was ultimately successful. The browser agent can now proceed to the next step of locating the search bar on Amazon's homepage.\n\nThe plan should continue with Step 2 as we have successfully completed Step 1.\nTerminate: False\nFinal Response: None", + "name": "critique_agent" + } +] \ No newline at end of file diff --git a/ta-browser/1744571514482277604/task_20250413_191154/conversation.json b/ta-browser/1744571514482277604/task_20250413_191154/conversation.json new file mode 100644 index 0000000..0637a08 --- /dev/null +++ b/ta-browser/1744571514482277604/task_20250413_191154/conversation.json @@ -0,0 +1 @@ +[] \ No newline at end of file diff --git a/ta-browser/1744572014568045555/task_20250413_192014/conversation.json b/ta-browser/1744572014568045555/task_20250413_192014/conversation.json new file mode 100644 index 0000000..d214df6 --- /dev/null +++ b/ta-browser/1744572014568045555/task_20250413_192014/conversation.json @@ -0,0 +1,49 @@ +[ + { + "role": "user", + "content": "Inputs :\nUser Query : Go to Amazon website and search for \"rtx 4060\". Extract the product listings including prices, models, and ratings for RTX 4060 graphics cards.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" + }, + { + "role": "assistant", + "content": "Plan: 1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\nNext Step: Open Amazon's website via direct URL: https://www.amazon.com", + "name": "planner_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" current_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" " + }, + { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", + "type": "function", + "function": { + "name": "open_url_tool", + "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", + "name": "open_url_tool", + "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" + }, + { + "role": "assistant", + "content": "I encountered an error while trying to open Amazon's website. The page failed to load successfully.", + "name": "browser_nav_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" next_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" tool_response=\"I encountered an error while trying to open Amazon's website. The page failed to load successfully.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" + }, + { + "role": "assistant", + "content": "Feedback: The original plan involves opening Amazon's website and then performing a series of actions to search for 'RTX 4060' and extract product details. The current step was to open Amazon's website using the direct URL. However, the browser agent encountered an error while trying to load the page. The error message indicates that the execution context was destroyed, likely due to a navigation issue. This means that the page could not be successfully loaded, and thus we cannot proceed with the search and extraction tasks. Since this is a critical failure at the very first step, it's not possible to continue with the plan.\nTerminate: True\nFinal Response: The task encountered an error while trying to open Amazon's website, which prevented further actions. The error was related to the execution context being destroyed, likely due to a navigation issue. This means the page could not be loaded successfully, and as a result, we cannot proceed with searching for RTX 4060 graphics cards or extracting any product information. I recommend retrying the task or checking the network connection and URL for any issues.", + "name": "critique_agent" + } +] \ No newline at end of file diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 9dbf16e..6a2aaef 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -244,27 +244,32 @@ def log_token_usage(self, agent_type: str, usage: Usage, step: Optional[int] = N """ ) - async def async_init(self, job_id: str, start_url: str = "https://google.com"): + async def async_init(self, job_id: str, start_url: str = "https://google.com", model_preference: str = "Anthropic"): """Initialize a new session context with improved error handling""" try: logger.info("Initializing browser session", extra={ "job_id": job_id, - "start_url": start_url + "start_url": start_url, + "model_preference": model_preference }) - - # 1. Initialize job_id and validate + print(f"[ORCHESTRATOR] Starting async_init with model_preference: {model_preference}") + + # Store and validate parameters + self.job_id = job_id + if not job_id: - raise ValueError("job_id is required for initialization") - self.job_id = str(job_id) - logger.debug(f"job_id: {self.job_id}") + raise ValueError("job_id is required") - # 2. Initialize conversation storage + # 1. Create conversation_storage with job_id + self.conversation_storage = ConversationStorage(job_id) + + # 2. Initialize browser manager try: - self.conversation_storage = ConversationStorage(job_id=self.job_id) - except Exception as storage_error: - raise RuntimeError(f"Failed to initialize conversation storage: {str(storage_error)}") from storage_error - - # 3. Set and validate URL + await self.initialize_browser_manager(start_url) + except Exception as browser_error: + raise RuntimeError(f"Failed to initialize browser: {str(browser_error)}") from browser_error + + # 3. Set current URL and domain try: self.current_url = start_url self.current_domain = extract_domain(self.current_url) @@ -273,23 +278,18 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com"): except InvalidURLError as url_error: raise ValueError(f"Invalid URL provided: {str(url_error)}") from url_error - # 4. Initialize browser manager with start_url - if not self.browser_manager: - try: - self.browser_manager = await self.initialize_browser_manager(start_url=start_url) - if not self.browser_manager: - raise RuntimeError("Browser manager initialization failed") - except Exception as browser_error: - raise RuntimeError(f"Failed to initialize browser manager: {str(browser_error)}") from browser_error - - # 5. Initialize client and agents + # 4. Initialize client and agents try: - + print(f"[ORCHESTRATOR] Calling initialize_client with model_preference: {model_preference}") from core.utils.init_client import initialize_client - self.client, model_instance = await initialize_client() + self.client, model_instance = await initialize_client(model_preference) + print(f"[ORCHESTRATOR] Client initialized successfully, now initializing agents") self.initialize_agents(model_instance) + print(f"[ORCHESTRATOR] Agents initialized successfully with model_preference: {model_preference}") except Exception as agent_error: - raise RuntimeError(f"Failed to initialize client and agents: {str(agent_error)}") from agent_error + error_msg = f"Failed to initialize client and agents: {str(agent_error)}" + print(f"[ORCHESTRATOR] ERROR: {error_msg}") + raise RuntimeError(error_msg) from agent_error self.async_init_done = True logger.debug("Async initialization completed successfully") diff --git a/ta-browser/core/server/models/web.py b/ta-browser/core/server/models/web.py index bf9e9fc..ca397f0 100644 --- a/ta-browser/core/server/models/web.py +++ b/ta-browser/core/server/models/web.py @@ -7,6 +7,7 @@ class StreamRequestModel(BaseModel): cmd: str = Field(..., description="Command to execute") url: str = Field("https://google.com", description="URL to navigate to") critique_disabled: bool = Field(False, description="Whether to disable critique") + model_preference: str = Field("Anthropic", description="Model provider preference (Anthropic or OpenAI)") @validator('url') def validate_and_format_url(cls, v): diff --git a/ta-browser/core/server/routes/web.py b/ta-browser/core/server/routes/web.py index a5cee6d..0f0ee09 100644 --- a/ta-browser/core/server/routes/web.py +++ b/ta-browser/core/server/routes/web.py @@ -115,6 +115,7 @@ async def stream_session( and streams back real-time updates as the command is executed. """ print(f"[{time.time()}] Stream route: Starting processing") + print(f"[STREAM] Received request with model_preference: {request.model_preference}") session_tracker = SessionTracker() @@ -124,8 +125,12 @@ async def stream_session( # Initialize the session logger.debug(f"Initializing stream session with ID {session_id}") + print(f"[STREAM] Initializing session with model_preference: {request.model_preference}") session_info = await session_tracker.initialize_session( - request.url, request.critique_disabled, session_id + request.url, + request.critique_disabled, + session_id, + request.model_preference ) session_context = session_tracker.active_sessions.get(session_id) diff --git a/ta-browser/core/server/utils/session_tracker.py b/ta-browser/core/server/utils/session_tracker.py index acf7168..fb3fa13 100644 --- a/ta-browser/core/server/utils/session_tracker.py +++ b/ta-browser/core/server/utils/session_tracker.py @@ -28,17 +28,19 @@ def get_active_sessions_status(self) -> dict: "sessions": list(self.active_sessions.keys()) } - async def initialize_session(self, start_url: str, no_crit: bool, session_id: str) -> Dict[str, Any]: + async def initialize_session(self, start_url: str, no_crit: bool, session_id: str, model_preference: str = "Anthropic") -> Dict[str, Any]: """Initialize a new session context""" logger.debug(f"Starting session initialization with URL: {start_url}") + print(f"[SESSION_TRACKER] Initializing session with model_preference: {model_preference}") logger.set_job_id(session_id) orchestrator = None try: orchestrator = Orchestrator(input_mode="API", no_crit=no_crit) - await orchestrator.async_init(job_id=session_id, start_url=start_url) - logger.debug(f"Orchestrator async_init completed with params: job_id={session_id}, start_url={start_url}") + print(f"[SESSION_TRACKER] Passing model_preference: {model_preference} to orchestrator.async_init") + await orchestrator.async_init(job_id=session_id, start_url=start_url, model_preference=model_preference) + logger.debug(f"Orchestrator async_init completed with params: job_id={session_id}, start_url={start_url}, model_preference={model_preference}") notification_queue = Queue() orchestrator.notification_queue = notification_queue @@ -49,7 +51,8 @@ async def initialize_session(self, start_url: str, no_crit: bool, session_id: st "notification_queue": notification_queue, "start_time": datetime.now(), "current_url": start_url, - "include_screenshot": False + "include_screenshot": False, + "model_preference": model_preference } self.add_active_session(session_id, session_context) diff --git a/ta-browser/core/utils/init_client.py b/ta-browser/core/utils/init_client.py index 7d170c6..5a423b9 100644 --- a/ta-browser/core/utils/init_client.py +++ b/ta-browser/core/utils/init_client.py @@ -1,44 +1,88 @@ import os from core.utils.anthropic_client import create_client_with_retry as create_anthropic_client, AsyncAnthropic from core.utils.logger import Logger +from core.utils.openai_client import get_client logger = Logger() -async def initialize_client(): +async def initialize_client(model_preference: str = "Anthropic"): """ Initialize and return the Anthropic client and model instance + Args: + model_preference (str): The model provider to use ("Anthropic" or "OpenAI") + Returns: tuple: (client_instance, model_instance) """ try: - # Get API key from environment variable - api_key = os.getenv("ANTHROPIC_API_KEY") - if not api_key: - logger.error("ANTHROPIC_API_KEY not found in environment variables") - raise ValueError("ANTHROPIC_API_KEY not found in environment variables") - - # Set model name - Claude 3.5 Sonnet - model_name = os.getenv("ANTHROPIC_MODEL_NAME") - - # Create client config - config = { - "api_key": api_key, - "model": model_name, - "max_retries": 3, - "timeout": 300.0 - } - - # Initialize client - client_instance = create_anthropic_client(AsyncAnthropic, config) - - # Create model instance - from pydantic_ai.models.anthropic import AnthropicModel - model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) - - logger.info(f"Anthropic client initialized successfully with model: {model_name}") - return client_instance, model_instance - + print(f"[INIT_CLIENT] *** Initializing client with model preference: {model_preference} ***") + logger.info(f"Initializing client with model preference: {model_preference}") + if model_preference == "Anthropic": + # Get API key from environment variable + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + logger.error("ANTHROPIC_API_KEY not found in environment variables") + raise ValueError("ANTHROPIC_API_KEY not found in environment variables") + + # Set model name - Claude 3.5 Sonnet + model_name = os.getenv("ANTHROPIC_MODEL_NAME") + print(f"[INIT_CLIENT] Using Anthropic model: {model_name}") + + # Create client config + config = { + "api_key": api_key, + "model": model_name, + "max_retries": 3, + "timeout": 300.0 + } + + # Initialize client + client_instance = create_anthropic_client(AsyncAnthropic, config) + + # Create model instance + from pydantic_ai.models.anthropic import AnthropicModel + model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) + + print(f"[INIT_CLIENT] Anthropic client initialized successfully with model: {model_name}") + logger.info(f"Anthropic client initialized successfully with model: {model_name}") + return client_instance, model_instance + elif model_preference == "OpenAI": + # Get API key from environment variable + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + logger.error("OPENAI_API_KEY not found in environment variables") + raise ValueError("OPENAI_API_KEY not found in environment variables") + + # Set model name - GPT-4o + model_name = os.getenv("OPENAI_MODEL_NAME") + print(f"[INIT_CLIENT] Using OpenAI model: {model_name}") + + # Create client config + config = { + "api_key": api_key, + "model": model_name, + "max_retries": 3, + "timeout": 300.0 + } + + # Initialize client + client_instance = get_client() + + # Create model instance + from pydantic_ai.models.openai import OpenAIModel + model_instance = OpenAIModel(model_name=model_name, openai_client=client_instance) + + print(f"[INIT_CLIENT] OpenAI client initialized successfully with model: {model_name}") + logger.info(f"OpenAI client initialized successfully with model: {model_name}") + return client_instance, model_instance + else: + error_msg = f"Invalid model preference: {model_preference}. Must be 'Anthropic' or 'OpenAI'" + print(f"[INIT_CLIENT] ERROR: {error_msg}") + raise ValueError(error_msg) + except Exception as e: - logger.error(f"Error initializing Anthropic client: {str(e)}", exc_info=True) + error_msg = f"Error initializing client: {str(e)}" + print(f"[INIT_CLIENT] CRITICAL ERROR: {error_msg}") + logger.error(error_msg, exc_info=True) raise \ No newline at end of file diff --git a/ta-browser/core/utils/openai_client.py b/ta-browser/core/utils/openai_client.py index 80670e3..5c74992 100644 --- a/ta-browser/core/utils/openai_client.py +++ b/ta-browser/core/utils/openai_client.py @@ -28,57 +28,15 @@ def validate_model(model: str) -> bool: return True @staticmethod def get_text_config() -> Dict: - model = get_env_var("AGENTIC_BROWSER_TEXT_MODEL") - if not OpenAIConfig.validate_model(model): - raise ModelValidationError( - f"Invalid model: {model}. Must match one of the patterns: " - f"{', '.join(OpenAIConfig.VALID_MODEL_PATTERNS)}" - ) - return { - "api_key": get_env_var("AGENTIC_BROWSER_TEXT_API_KEY"), - "base_url": get_env_var("AGENTIC_BROWSER_TEXT_BASE_URL"), - "model": model, + "api_key": get_env_var("OPENAI_API_KEY"), + "base_url": "https://api.openai.com/v1", + "model": get_env_var("OPENAI_MODEL_NAME"), "max_retries": 3, "timeout": 300.0 } - @staticmethod - def get_ss_config() -> Dict: - model = get_env_var("AGENTIC_BROWSER_SS_MODEL") - if not OpenAIConfig.validate_model(model): - raise ModelValidationError( - f"Invalid model: {model}. Must match one of the patterns: " - f"{', '.join(OpenAIConfig.VALID_MODEL_PATTERNS)}" - ) - - return { - "api_key": get_env_var("AGENTIC_BROWSER_SS_API_KEY"), - "base_url": get_env_var("AGENTIC_BROWSER_SS_BASE_URL"), - "model": model, - "max_retries": 3, - "timeout": 300.0 - } -async def validate_models(client: AsyncOpenAI) -> bool: - """Validate that configured models are available""" - try: - available_models = await client.models.list() - available_model_ids = [model.id for model in available_models.data] - - text_model = get_text_model() - ss_model = get_ss_model() - - if text_model not in available_model_ids: - raise ModelValidationError(f"Text model '{text_model}' not available. Available models: {', '.join(available_model_ids)}") - - if ss_model not in available_model_ids: - raise ModelValidationError(f"Screenshot model '{ss_model}' not available. Available models: {', '.join(available_model_ids)}") - - return True - except Exception as e: - logger.error(f"Model validation failed: {str(e)}") - return False def create_client_with_retry(client_class, config: dict): """Create an OpenAI client with proper error handling""" @@ -102,26 +60,9 @@ def get_client(): config = OpenAIConfig.get_text_config() return create_client_with_retry(AsyncOpenAI, config) -def get_ss_client(): - """Get OpenAI client for screenshot analysis""" - config = OpenAIConfig.get_ss_config() - return create_client_with_retry(OpenAI, config) - -def get_text_model() -> str: - """Get model name for text analysis""" - return OpenAIConfig.get_text_config()["model"] - -def get_ss_model() -> str: - """Get model name for screenshot analysis""" - return OpenAIConfig.get_ss_config()["model"] - # Example usage async def initialize_and_validate(): """Initialize client and validate configuration""" client = get_client() - # Validate models - if not await validate_models(client): - raise ModelValidationError("Failed to validate models. Please check your configuration.") - return client \ No newline at end of file From aebd6eafcd8ba704f79cb794fef09c4aca573b2e Mon Sep 17 00:00:00 2001 From: aryan Date: Mon, 14 Apr 2025 16:47:23 +0530 Subject: [PATCH 02/24] fix(conversation_handling): Remove obsolete conversation files and enhance error handling in orchestrator. - Deleted unused conversation JSON files - Improved error handling during conversation storage initialization in the orchestrator to provide clearer feedback on failures. --- .../task_20250413_190916/conversation.json | 79 ------------------- .../task_20250413_191154/conversation.json | 1 - .../task_20250413_192014/conversation.json | 49 ------------ ta-browser/core/orchestrator.py | 5 +- ta-browser/core/utils/openai_msg_parser.py | 1 - 5 files changed, 4 insertions(+), 131 deletions(-) delete mode 100644 ta-browser/1744571356739037821/task_20250413_190916/conversation.json delete mode 100644 ta-browser/1744571514482277604/task_20250413_191154/conversation.json delete mode 100644 ta-browser/1744572014568045555/task_20250413_192014/conversation.json diff --git a/ta-browser/1744571356739037821/task_20250413_190916/conversation.json b/ta-browser/1744571356739037821/task_20250413_190916/conversation.json deleted file mode 100644 index e9976ca..0000000 --- a/ta-browser/1744571356739037821/task_20250413_190916/conversation.json +++ /dev/null @@ -1,79 +0,0 @@ -[ - { - "role": "user", - "content": "Inputs :\nUser Query : Go to Amazon.com and search for RTX 4060 graphics cards. Return the search results including prices and available models.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" - }, - { - "role": "assistant", - "content": "Plan: 1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\nNext Step: Navigate to Amazon's website (https://www.amazon.com)", - "name": "planner_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" current_step=\"Navigate to Amazon's website (https://www.amazon.com)\" " - }, - { - "role": "assistant", - "content": "I'll help navigate to Amazon's website using the open_url_tool.", - "name": "browser_nav_agent" - }, - { - "role": "assistant", - "content": null, - "tool_calls": [ - { - "id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", - "type": "function", - "function": { - "name": "open_url_tool", - "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" - } - } - ] - }, - { - "role": "tool", - "tool_call_id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", - "name": "open_url_tool", - "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" - }, - { - "role": "assistant", - "content": "I attempted to navigate to Amazon's website but encountered an error. This is common with Amazon as they have strict bot detection measures. Let me verify if we landed on any page by checking the current URL.", - "name": "browser_nav_agent" - }, - { - "role": "assistant", - "content": null, - "tool_calls": [ - { - "id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", - "type": "function", - "function": { - "name": "get_url", - "arguments": "{\"raw_args\": \"ArgsDict(args_dict={})\"}" - } - } - ] - }, - { - "role": "tool", - "tool_call_id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", - "name": "get_url", - "content": "Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more." - }, - { - "role": "assistant", - "content": "Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.", - "name": "browser_nav_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" next_step=\"Navigate to Amazon's website (https://www.amazon.com)\" tool_response=\"Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\nTool Call: get_url\nArguments: {}\nResponse: Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more.\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" - }, - { - "role": "assistant", - "content": "Feedback: Original Plan:\n1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\n\nCurrent Progress:\nCurrently on Step 1: Navigate to Amazon's website (https://www.amazon.com)\n\nAnalysis and Feedback:\nThe browser agent has successfully completed the first step of navigating to Amazon.com. This is confirmed by:\n1. The get_url tool response showing we are at \"https://www.amazon.com/\"\n2. The page title verification showing \"Amazon.com. Spend less. Smile more.\"\n\nDespite an initial execution context error (which is common during page loads), the navigation was ultimately successful. The browser agent can now proceed to the next step of locating the search bar on Amazon's homepage.\n\nThe plan should continue with Step 2 as we have successfully completed Step 1.\nTerminate: False\nFinal Response: None", - "name": "critique_agent" - } -] \ No newline at end of file diff --git a/ta-browser/1744571514482277604/task_20250413_191154/conversation.json b/ta-browser/1744571514482277604/task_20250413_191154/conversation.json deleted file mode 100644 index 0637a08..0000000 --- a/ta-browser/1744571514482277604/task_20250413_191154/conversation.json +++ /dev/null @@ -1 +0,0 @@ -[] \ No newline at end of file diff --git a/ta-browser/1744572014568045555/task_20250413_192014/conversation.json b/ta-browser/1744572014568045555/task_20250413_192014/conversation.json deleted file mode 100644 index d214df6..0000000 --- a/ta-browser/1744572014568045555/task_20250413_192014/conversation.json +++ /dev/null @@ -1,49 +0,0 @@ -[ - { - "role": "user", - "content": "Inputs :\nUser Query : Go to Amazon website and search for \"rtx 4060\". Extract the product listings including prices, models, and ratings for RTX 4060 graphics cards.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" - }, - { - "role": "assistant", - "content": "Plan: 1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\nNext Step: Open Amazon's website via direct URL: https://www.amazon.com", - "name": "planner_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" current_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" " - }, - { - "role": "assistant", - "content": null, - "tool_calls": [ - { - "id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", - "type": "function", - "function": { - "name": "open_url_tool", - "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" - } - } - ] - }, - { - "role": "tool", - "tool_call_id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", - "name": "open_url_tool", - "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" - }, - { - "role": "assistant", - "content": "I encountered an error while trying to open Amazon's website. The page failed to load successfully.", - "name": "browser_nav_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" next_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" tool_response=\"I encountered an error while trying to open Amazon's website. The page failed to load successfully.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" - }, - { - "role": "assistant", - "content": "Feedback: The original plan involves opening Amazon's website and then performing a series of actions to search for 'RTX 4060' and extract product details. The current step was to open Amazon's website using the direct URL. However, the browser agent encountered an error while trying to load the page. The error message indicates that the execution context was destroyed, likely due to a navigation issue. This means that the page could not be successfully loaded, and thus we cannot proceed with the search and extraction tasks. Since this is a critical failure at the very first step, it's not possible to continue with the plan.\nTerminate: True\nFinal Response: The task encountered an error while trying to open Amazon's website, which prevented further actions. The error was related to the execution context being destroyed, likely due to a navigation issue. This means the page could not be loaded successfully, and as a result, we cannot proceed with searching for RTX 4060 graphics cards or extracting any product information. I recommend retrying the task or checking the network connection and URL for any issues.", - "name": "critique_agent" - } -] \ No newline at end of file diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 6a2aaef..69df1d7 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -261,7 +261,10 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com", m raise ValueError("job_id is required") # 1. Create conversation_storage with job_id - self.conversation_storage = ConversationStorage(job_id) + try: + self.conversation_storage = ConversationStorage(job_id=self.job_id) + except Exception as storage_error: + raise RuntimeError(f"Failed to initialize conversation storage: {str(storage_error)}") from storage_error # 2. Initialize browser manager try: diff --git a/ta-browser/core/utils/openai_msg_parser.py b/ta-browser/core/utils/openai_msg_parser.py index af57cdc..1458125 100644 --- a/ta-browser/core/utils/openai_msg_parser.py +++ b/ta-browser/core/utils/openai_msg_parser.py @@ -646,7 +646,6 @@ def get_conversation_history(self) -> List[ChatCompletionMessageParam]: return self.conversation_history -# In openai_msg_parser.py (or wherever ConversationStorage is defined) from datetime import datetime import os import json From 61f257d4f462ded07be195590356a3bcb4584b58 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Wed, 16 Apr 2025 15:52:56 +0530 Subject: [PATCH 03/24] feat: agent reinitialization everytime the model_preference changes --- cortex_on/agents/code_agent.py | 23 +++++---- cortex_on/agents/hello.py | 65 ++++++++++++++++++++++++++ cortex_on/agents/orchestrator_agent.py | 54 +++++++++++++++++---- cortex_on/agents/planner_agent.py | 29 +++++++----- cortex_on/instructor.py | 16 +++++-- cortex_on/main.py | 32 +++++++++++-- cortex_on/utils/ant_client.py | 24 ++++++++++ ta-browser/core/server/routes/web.py | 14 +++++- 8 files changed, 221 insertions(+), 36 deletions(-) create mode 100644 cortex_on/agents/hello.py diff --git a/cortex_on/agents/code_agent.py b/cortex_on/agents/code_agent.py index 43fa047..8580ab6 100644 --- a/cortex_on/agents/code_agent.py +++ b/cortex_on/agents/code_agent.py @@ -15,7 +15,7 @@ from pydantic_ai.models.anthropic import AnthropicModel # Local application imports -from utils.ant_client import get_client +from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client from utils.stream_response_format import StreamResponse load_dotenv() @@ -25,6 +25,7 @@ class CoderAgentDeps: websocket: Optional[WebSocket] = None stream_output: Optional[StreamResponse] = None + model_preference: str = "Anthropic" # Constants ALLOWED_COMMANDS = { @@ -236,20 +237,26 @@ async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> N stream_output_json = json.dumps(asdict(ctx.deps.stream_output)) logfire.debug("WebSocket message sent: {stream_output_json}", stream_output_json=stream_output_json) -# Initialize the model -model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() -) + # Initialize the agent -coder_agent = Agent( + +def coder_agent(model_preference: str = "Anthropic"): + if model_preference == "Anthropic": + model = get_anthropic_model_instance() + elif model_preference == "OpenAI": + model = get_openai_model_instance() + else: + raise ValueError(f"Unknown model_preference: {model_preference}") + print(f"[CODER_INIT] Creating coder agent with model: {model}") + coder_agent = Agent( model=model, name="Coder Agent", result_type=CoderResult, deps_type=CoderAgentDeps, system_prompt=coder_system_message -) + ) + return coder_agent @coder_agent.tool async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: diff --git a/cortex_on/agents/hello.py b/cortex_on/agents/hello.py new file mode 100644 index 0000000..2c453cc --- /dev/null +++ b/cortex_on/agents/hello.py @@ -0,0 +1,65 @@ + +@dataclass +class orchestrator_deps: + websocket: Optional[WebSocket] = None + stream_output: Optional[StreamResponse] = None + # Add a collection to track agent-specific streams + agent_responses: Optional[List[StreamResponse]] = None + model_preference: str = "Anthropic" + +async def get_model_preference() -> str: + """Fetch model preference from ta_browser API""" + try: + async with aiohttp.ClientSession() as session: + async with session.get("http://localhost:8000/api/v1/model/preference") as response: + if response.status == 200: + data = await response.json() + return data.get("model_preference", "Anthropic") + else: + logfire.error(f"Failed to get model preference: {response.status}") + return "Anthropic" + except Exception as e: + logfire.error(f"Error getting model preference: {str(e)}") + return "Anthropic" + +async def initialize_orchestrator_agent(): + """Initialize the orchestrator agent with model preference from ta_browser""" + model_preference = await get_model_preference() + model = AnthropicModel( + model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), + anthropic_client=get_client() + ) + + print(f"Orchestrator agent model initialized with MODEL PREFERENCE: {model_preference}") + + return Agent( + model=model, + name="Orchestrator Agent", + system_prompt=orchestrator_system_prompt, + deps_type=orchestrator_deps + ) + +# Initialize the agent +orchestrator_agent = None + +def get_orchestrator_agent(): + """Get the orchestrator agent, initializing it if necessary""" + global orchestrator_agent + if orchestrator_agent is None: + raise RuntimeError("Orchestrator agent not initialized. Call initialize_orchestrator_agent() first.") + return orchestrator_agent + +@get_orchestrator_agent().tool +async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Plans the task and assigns it to the appropriate agents""" + try: + logfire.info(f"Planning task: {task}") + + # Create a new StreamResponse for Planner Agent + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 284d927..fb92b3c 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -14,7 +14,9 @@ from utils.stream_response_format import StreamResponse from agents.planner_agent import planner_agent, update_todo_status from agents.code_agent import coder_agent, CoderAgentDeps -from utils.ant_client import get_client +from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client + +print("[ORCH_INIT] Starting orchestrator agent initialization") @dataclass class orchestrator_deps: @@ -24,6 +26,7 @@ class orchestrator_deps: agent_responses: Optional[List[StreamResponse]] = None model_preference: str = "Anthropic" + orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. [AGENT CAPABILITIES] @@ -159,17 +162,39 @@ class orchestrator_deps: - Format: "Task description (agent_name)" """ -model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() -) +# print("[ORCH_INIT] Getting Anthropic client...") +# client = get_client() +# print(f"[ORCH_INIT] Anthropic client obtained: {client}") + +print("[ORCH_INIT] Initializing Anthropic model...") +model = get_anthropic_model_instance() +print(f"[ORCH_INIT] Anthropic model initialized: {model}") + +print("Orchestrator agent model initialized without MODEL PREFERENCE") + +print("[ORCH_INIT] Creating orchestrator agent...") orchestrator_agent = Agent( model=model, name="Orchestrator Agent", system_prompt=orchestrator_system_prompt, deps_type=orchestrator_deps ) +print("[ORCH_INIT] Orchestrator agent created successfully") +def orchestrator_agent(model_preference: str): + if model_preference == "Anthropic": + model = get_anthropic_model_instance() + elif model_preference == "OpenAI": + model = get_openai_model_instance() + else: + raise ValueError(f"Unknown model_preference: {model_preference}") + print(f"[ORCH_INIT] Creating orchestrator agent with model: {model}") + return Agent( + model=model, + name="Orchestrator Agent", + system_prompt=orchestrator_system_prompt, + deps_type=orchestrator_deps + ) @orchestrator_agent.tool async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: @@ -197,7 +222,10 @@ async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) # Run planner agent - planner_response = await planner_agent.run(user_prompt=task) + agent = planner_agent(model_preference=ctx.deps.model_preference) + planner_response = await agent.run(user_prompt=task) + + logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") # Update planner stream with results plan_text = planner_response.data.plan @@ -253,14 +281,18 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: # Create deps with the new stream_output deps_for_coder_agent = CoderAgentDeps( websocket=ctx.deps.websocket, - stream_output=coder_stream_output + stream_output=coder_stream_output, + model_preference = ctx.deps.model_preference ) # Run coder agent - coder_response = await coder_agent.run( + agent = coder_agent(model_preference=ctx.deps.model_preference) + coder_response = await agent.run( user_prompt=task, deps=deps_for_coder_agent ) + logfire.info(f"coder_response: {coder_response}") + logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") # Extract response data response_data = coder_response.data.content @@ -313,6 +345,7 @@ async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: api_url="http://localhost:8000/api/v1/web/stream", model_preference=ctx.deps.model_preference ) + logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") # Run WebSurfer with its own stream_output success, message, messages = await web_surfer_agent.generate_reply( @@ -442,7 +475,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas # We'll directly call planner_agent.run() to create a new plan first plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - plan_response = await planner_agent.run(user_prompt=plan_prompt) + agent = planner_agent(model_preference=ctx.deps.model_preference) + plan_response = await agent.run(user_prompt=plan_prompt) current_content = plan_response.data.plan else: # Read existing todo.md @@ -488,7 +522,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.error(error_msg, exc_info=True) planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = a500 + planner_stream_output.status_code = 500 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) return f"Failed to update the plan: {error_msg}" diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 897c22f..53592ae 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -10,7 +10,7 @@ from pydantic_ai.models.anthropic import AnthropicModel # Local application imports -from utils.ant_client import get_client +from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client @@ -176,17 +176,24 @@ class PlannerResult(BaseModel): plan: str = Field(description="The generated or updated plan in string format - this should be the complete plan text") -model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() -) -planner_agent = Agent( - model=model, - name="Planner Agent", - result_type=PlannerResult, - system_prompt=planner_prompt -) + + +def planner_agent(model_preference: str = "Anthropic"): + if model_preference == "Anthropic": + model = get_anthropic_model_instance() + elif model_preference == "OpenAI": + model = get_openai_model_instance() + else: + raise ValueError(f"Unknown model_preference: {model_preference}") + print(f"[PLANNER_INIT] Creating planner agent with model: {model}") + planner_agent = Agent( + model=model, + name="Planner Agent", + result_type=PlannerResult, + system_prompt=planner_prompt + ) + return planner_agent @planner_agent.tool_plain async def update_todo_status(task_description: str) -> str: diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index f2688b4..d92cf4e 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -25,8 +25,10 @@ load_dotenv() - - +print("[INIT] Loading agents...") +print(f"[INIT] Planner agent loaded: {planner_agent}") +print(f"[INIT] Coder agent loaded: {coder_agent}") +print(f"[INIT] Orchestrator agent loaded: {get_orchestrator_agent}") class DateTimeEncoder(json.JSONEncoder): """Custom JSON encoder that can handle datetime objects""" @@ -39,11 +41,13 @@ def default(self, obj): # Main Orchestrator Class class SystemInstructor: def __init__(self, model_preference: str = "Anthropic"): + print(f"[INIT] Initializing SystemInstructor with model_preference: {model_preference}") self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None self.orchestrator_response: List[StreamResponse] = [] self.model_preference = model_preference self._setup_logging() + print("[INIT] SystemInstructor initialization complete") def _setup_logging(self) -> None: """Configure logging with proper formatting""" @@ -91,10 +95,16 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: stream_output.steps.append("Agents initialized successfully") await self._safe_websocket_send(stream_output) - orchestrator_response = await orchestrator_agent.run( + print("Orchestrator agent running") + + agent = orchestrator_agent(self.model_preference) + orchestrator_response = await agent.run( user_prompt=task, deps=deps_for_orchestrator ) + + logfire.info(f"Orchestrator Agent using model type: {self.model_preference}") + stream_output.output = orchestrator_response.data stream_output.status_code = 200 logfire.debug(f"Orchestrator response: {orchestrator_response.data}") diff --git a/cortex_on/main.py b/cortex_on/main.py index e1adddd..4bab62f 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -13,12 +13,25 @@ # Default model preference is Anthropic MODEL_PREFERENCE = "Anthropic" +# Global instructor instance +instructor = None + @asynccontextmanager async def lifespan(app: FastAPI): # Set default model preference at startup app.state.model_preference = MODEL_PREFERENCE print(f"[STARTUP] Setting default model preference to: {MODEL_PREFERENCE}") + + # Initialize the instructor + global instructor + instructor = SystemInstructor(model_preference=MODEL_PREFERENCE) + print("[STARTUP] Instructor initialized") + yield + + # Cleanup + if instructor: + await instructor.shutdown() app: FastAPI = FastAPI(lifespan=lifespan) @@ -42,21 +55,34 @@ async def get_model_preference() -> str: @app.get("/set_model_preference") async def set_model_preference(model: str): """ - Set the model preference (Anthropic or OpenAI) + Set the model preference (Anthropic or OpenAI) and reinitialize the instructor """ if model not in ["Anthropic", "OpenAI"]: print(f"[ERROR] Invalid model preference attempted: {model}") return {"error": "Model must be 'Anthropic' or 'OpenAI'"} + print(f"[SET] Changing model preference from {app.state.model_preference} to {model}") app.state.model_preference = model + + # Reinitialize the instructor with new model preference + global instructor + if instructor: + await instructor.shutdown() + instructor = SystemInstructor(model_preference=model) + print(f"[SET] Instructor reinitialized with model preference: {model}") + return {"message": f"Model preference set to {model}"} async def generate_response(task: str, websocket: Optional[WebSocket] = None, model_preference: str = None): if model_preference is None: model_preference = app.state.model_preference print(f"[GENERATE] Using model preference: {model_preference} for task: {task[:30]}...") - orchestrator: SystemInstructor = SystemInstructor(model_preference) - return await orchestrator.run(task, websocket) + + global instructor + if not instructor: + instructor = SystemInstructor(model_preference=model_preference) + + return await instructor.run(task, websocket) @app.get("/agent/chat") async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List: diff --git a/cortex_on/utils/ant_client.py b/cortex_on/utils/ant_client.py index 924f0c6..c04732c 100644 --- a/cortex_on/utils/ant_client.py +++ b/cortex_on/utils/ant_client.py @@ -1,5 +1,6 @@ from pydantic_ai.models.anthropic import AnthropicModel from anthropic import AsyncAnthropic +from openai import AsyncOpenAI import os from dotenv import load_dotenv @@ -12,3 +13,26 @@ def get_client(): max_retries=3, timeout=10000) return client + + + +def get_openai_client(): + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key, + max_retries=3, + timeout=10000) + return client + +def get_openai_model_instance(): + model_name = os.getenv("OPENAI_MODEL_NAME") + # Create model instance + from pydantic_ai.models.openai import OpenAIModel + model_instance = OpenAIModel(model_name=model_name, openai_client=get_openai_client()) + return model_instance + +def get_anthropic_model_instance(): + model_name = os.getenv("ANTHROPIC_MODEL_NAME") + # Create model instance + from pydantic_ai.models.anthropic import AnthropicModel + model_instance = AnthropicModel(model_name=model_name, anthropic_client=get_client()) + return model_instance diff --git a/ta-browser/core/server/routes/web.py b/ta-browser/core/server/routes/web.py index 0f0ee09..8dabe34 100644 --- a/ta-browser/core/server/routes/web.py +++ b/ta-browser/core/server/routes/web.py @@ -1,17 +1,19 @@ import asyncio -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, WebSocket from fastapi.responses import StreamingResponse from typing import AsyncGenerator import json import time from datetime import datetime from queue import Empty +import os from core.server.models.web import StreamRequestModel, StreamResponseModel from core.server.constants import GLOBAL_TIMEOUT from core.server.utils.timeout import timeout from core.server.utils.session_tracker import SessionTracker from core.utils.logger import Logger +from core.utils.init_client import initialize_client logger = Logger() @@ -150,4 +152,14 @@ async def stream_session( except Exception as e: logger.error(f"Error in /stream: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@router.get("/api/v1/model/preference") +async def get_model_preference(): + """Get the current model preference""" + try: + # Get model preference from environment or use default + model_preference = os.getenv("MODEL_PREFERENCE", "Anthropic") + return {"model_preference": model_preference} + except Exception as e: raise HTTPException(status_code=500, detail=str(e)) \ No newline at end of file From ba2083caa2543f8eed43a7303ead7ec32053a4d3 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Wed, 16 Apr 2025 17:25:23 +0530 Subject: [PATCH 04/24] feat: agent reinitialization everytime the model_preference changes --- cortex_on/agents/code_agent.py | 296 +++++------ cortex_on/agents/orchestrator_agent.py | 670 +++++++++++++------------ cortex_on/agents/planner_agent.py | 296 +++++------ cortex_on/main.py | 6 +- 4 files changed, 641 insertions(+), 627 deletions(-) diff --git a/cortex_on/agents/code_agent.py b/cortex_on/agents/code_agent.py index 8580ab6..d1eb785 100644 --- a/cortex_on/agents/code_agent.py +++ b/cortex_on/agents/code_agent.py @@ -241,7 +241,7 @@ async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> N # Initialize the agent -def coder_agent(model_preference: str = "Anthropic"): +async def coder_agent(model_preference: str = "Anthropic") -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -256,167 +256,169 @@ def coder_agent(model_preference: str = "Anthropic"): deps_type=CoderAgentDeps, system_prompt=coder_system_message ) - return coder_agent - -@coder_agent.tool -async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: - """ - Executes a shell command within a restricted directory and returns the output. - This consolidated tool handles terminal commands and file operations. - """ - try: - # Extract base command for security checks and messaging - base_command = command.split()[0] if command.split() else "" - - # Send operation description message - operation_message = get_high_level_operation_message(command, base_command) - await send_stream_update(ctx, operation_message) - - logfire.info("Executing shell command: {command}", command=command) - - # Setup restricted directory - base_dir = os.path.abspath(os.path.dirname(__file__)) - restricted_dir = os.path.join(base_dir, "code_files") - os.makedirs(restricted_dir, exist_ok=True) - - # Security check - if base_command not in ALLOWED_COMMANDS: - await send_stream_update(ctx, "Operation not permitted") - return f"Error: Command '{base_command}' is not allowed for security reasons." - - # Change to restricted directory for execution - original_dir = os.getcwd() - os.chdir(restricted_dir) - + + @coder_agent.tool + async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: + """ + Executes a shell command within a restricted directory and returns the output. + This consolidated tool handles terminal commands and file operations. + """ try: - # Handle echo with redirection (file writing) - if ">" in command and base_command == "echo": - file_path = command.split(">", 1)[1].strip() - await send_stream_update(ctx, f"Writing content to {file_path}") - - # Parse command parts - parts = command.split(">", 1) - echo_cmd = parts[0].strip() - - # Extract content, removing quotes if present - content = echo_cmd[5:].strip() - if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): - content = content[1:-1] - - try: - with open(file_path, "w") as file: - file.write(content) - - await send_stream_update(ctx, f"File {file_path} created successfully") - return f"Successfully wrote to {file_path}" - except Exception as e: - error_msg = f"Error writing to file: {str(e)}" - await send_stream_update(ctx, f"Failed to create file {file_path}") - logfire.error(error_msg, exc_info=True) - return error_msg + # Extract base command for security checks and messaging + base_command = command.split()[0] if command.split() else "" - # Handle cat with here-document for multiline file writing - elif "<<" in command and base_command == "cat": - cmd_parts = command.split("<<", 1) - cat_part = cmd_parts[0].strip() - - # Extract filename for status message if possible - file_path = None - if ">" in cat_part: - file_path = cat_part.split(">", 1)[1].strip() - await send_stream_update(ctx, f"Creating file {file_path}") + # Send operation description message + operation_message = get_high_level_operation_message(command, base_command) + await send_stream_update(ctx, operation_message) + + logfire.info("Executing shell command: {command}", command=command) + + # Setup restricted directory + base_dir = os.path.abspath(os.path.dirname(__file__)) + restricted_dir = os.path.join(base_dir, "code_files") + os.makedirs(restricted_dir, exist_ok=True) + + # Security check + if base_command not in ALLOWED_COMMANDS: + await send_stream_update(ctx, "Operation not permitted") + return f"Error: Command '{base_command}' is not allowed for security reasons." + + # Change to restricted directory for execution + original_dir = os.getcwd() + os.chdir(restricted_dir) + + try: + # Handle echo with redirection (file writing) + if ">" in command and base_command == "echo": + file_path = command.split(">", 1)[1].strip() + await send_stream_update(ctx, f"Writing content to {file_path}") + + # Parse command parts + parts = command.split(">", 1) + echo_cmd = parts[0].strip() + + # Extract content, removing quotes if present + content = echo_cmd[5:].strip() + if (content.startswith('"') and content.endswith('"')) or \ + (content.startswith("'") and content.endswith("'")): + content = content[1:-1] + + try: + with open(file_path, "w") as file: + file.write(content) + + await send_stream_update(ctx, f"File {file_path} created successfully") + return f"Successfully wrote to {file_path}" + except Exception as e: + error_msg = f"Error writing to file: {str(e)}" + await send_stream_update(ctx, f"Failed to create file {file_path}") + logfire.error(error_msg, exc_info=True) + return error_msg - try: - # Parse heredoc parts - doc_part = cmd_parts[1].strip() + # Handle cat with here-document for multiline file writing + elif "<<" in command and base_command == "cat": + cmd_parts = command.split("<<", 1) + cat_part = cmd_parts[0].strip() - # Extract filename + # Extract filename for status message if possible + file_path = None if ">" in cat_part: file_path = cat_part.split(">", 1)[1].strip() - else: - await send_stream_update(ctx, "Invalid file operation") - return "Error: Invalid cat command format. Must include redirection." + await send_stream_update(ctx, f"Creating file {file_path}") - # Parse the heredoc content and delimiter - if "\n" in doc_part: - delimiter_and_content = doc_part.split("\n", 1) - delimiter = delimiter_and_content[0].strip("'").strip('"') - content = delimiter_and_content[1] + try: + # Parse heredoc parts + doc_part = cmd_parts[1].strip() - # Find the end delimiter and extract content - if f"\n{delimiter}" in content: - content = content.split(f"\n{delimiter}")[0] - - # Write to file - with open(file_path, "w") as file: - file.write(content) + # Extract filename + if ">" in cat_part: + file_path = cat_part.split(">", 1)[1].strip() + else: + await send_stream_update(ctx, "Invalid file operation") + return "Error: Invalid cat command format. Must include redirection." + + # Parse the heredoc content and delimiter + if "\n" in doc_part: + delimiter_and_content = doc_part.split("\n", 1) + delimiter = delimiter_and_content[0].strip("'").strip('"') + content = delimiter_and_content[1] - await send_stream_update(ctx, f"File {file_path} created successfully") - return f"Successfully wrote multiline content to {file_path}" + # Find the end delimiter and extract content + if f"\n{delimiter}" in content: + content = content.split(f"\n{delimiter}")[0] + + # Write to file + with open(file_path, "w") as file: + file.write(content) + + await send_stream_update(ctx, f"File {file_path} created successfully") + return f"Successfully wrote multiline content to {file_path}" + else: + await send_stream_update(ctx, "File content format error") + return "Error: End delimiter not found in heredoc" else: await send_stream_update(ctx, "File content format error") - return "Error: End delimiter not found in heredoc" - else: - await send_stream_update(ctx, "File content format error") - return "Error: Invalid heredoc format" - except Exception as e: - error_msg = f"Error processing cat with heredoc: {str(e)}" - file_path_str = file_path if file_path else 'file' - await send_stream_update(ctx, f"Failed to create file {file_path_str}") - logfire.error(error_msg, exc_info=True) - return error_msg - - # Execute standard commands - else: - # Send execution message - execution_msg = get_high_level_execution_message(command, base_command) - await send_stream_update(ctx, execution_msg) + return "Error: Invalid heredoc format" + except Exception as e: + error_msg = f"Error processing cat with heredoc: {str(e)}" + file_path_str = file_path if file_path else 'file' + await send_stream_update(ctx, f"Failed to create file {file_path_str}") + logfire.error(error_msg, exc_info=True) + return error_msg - # Execute the command using subprocess - try: - args = shlex.split(command) - result = subprocess.run( - args, - shell=True, - capture_output=True, - text=True, - timeout=60, - ) + # Execute standard commands + else: + # Send execution message + execution_msg = get_high_level_execution_message(command, base_command) + await send_stream_update(ctx, execution_msg) - logfire.info(f"Command executed: {result.args}") + # Execute the command using subprocess + try: + args = shlex.split(command) + result = subprocess.run( + args, + shell=True, + capture_output=True, + text=True, + timeout=60, + ) + + logfire.info(f"Command executed: {result.args}") + + # Handle success + if result.returncode == 0: + success_msg = get_success_message(command, base_command) + await send_stream_update(ctx, success_msg) + logfire.info(f"Command executed successfully: {result.stdout}") + return result.stdout + + # Handle failure + else: + files = os.listdir('.') + error_msg = f"Command failed with error code {result.returncode}:\n{result.stderr}\n\nFiles in directory: {files}" + failure_msg = get_failure_message(command, base_command) + await send_stream_update(ctx, failure_msg) + return error_msg - # Handle success - if result.returncode == 0: - success_msg = get_success_message(command, base_command) - await send_stream_update(ctx, success_msg) - logfire.info(f"Command executed successfully: {result.stdout}") - return result.stdout + except subprocess.TimeoutExpired: + await send_stream_update(ctx, "Operation timed out") + return "Command execution timed out after 60 seconds" - # Handle failure - else: - files = os.listdir('.') - error_msg = f"Command failed with error code {result.returncode}:\n{result.stderr}\n\nFiles in directory: {files}" - failure_msg = get_failure_message(command, base_command) - await send_stream_update(ctx, failure_msg) + except Exception as e: + error_msg = f"Error executing command: {str(e)}" + await send_stream_update(ctx, "Operation failed") + logfire.error(error_msg, exc_info=True) return error_msg + + finally: + # Always return to the original directory + os.chdir(original_dir) - except subprocess.TimeoutExpired: - await send_stream_update(ctx, "Operation timed out") - return "Command execution timed out after 60 seconds" - - except Exception as e: - error_msg = f"Error executing command: {str(e)}" - await send_stream_update(ctx, "Operation failed") - logfire.error(error_msg, exc_info=True) - return error_msg + except Exception as e: + error_msg = f"Error executing command: {str(e)}" + await send_stream_update(ctx, "Operation failed") + logfire.error(error_msg, exc_info=True) + return error_msg - finally: - # Always return to the original directory - os.chdir(original_dir) - - except Exception as e: - error_msg = f"Error executing command: {str(e)}" - await send_stream_update(ctx, "Operation failed") - logfire.error(error_msg, exc_info=True) - return error_msg \ No newline at end of file + logfire.info("All tools initialized for coder agent") + return coder_agent \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index fb92b3c..4a5eec5 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -166,22 +166,22 @@ class orchestrator_deps: # client = get_client() # print(f"[ORCH_INIT] Anthropic client obtained: {client}") -print("[ORCH_INIT] Initializing Anthropic model...") -model = get_anthropic_model_instance() - -print(f"[ORCH_INIT] Anthropic model initialized: {model}") - -print("Orchestrator agent model initialized without MODEL PREFERENCE") - -print("[ORCH_INIT] Creating orchestrator agent...") -orchestrator_agent = Agent( - model=model, - name="Orchestrator Agent", - system_prompt=orchestrator_system_prompt, - deps_type=orchestrator_deps -) -print("[ORCH_INIT] Orchestrator agent created successfully") -def orchestrator_agent(model_preference: str): +# print("[ORCH_INIT] Initializing Anthropic model...") +# model = get_anthropic_model_instance() + +# print(f"[ORCH_INIT] Anthropic model initialized: {model}") + +# print("Orchestrator agent model initialized without MODEL PREFERENCE") + +# print("[ORCH_INIT] Creating orchestrator agent...") +# orchestrator_agent = Agent( +# model=model, +# name="Orchestrator Agent", +# system_prompt=orchestrator_system_prompt, +# deps_type=orchestrator_deps +# ) +# print("[ORCH_INIT] Orchestrator agent created successfully") +async def orchestrator_agent(model_preference: str) -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -189,354 +189,364 @@ def orchestrator_agent(model_preference: str): else: raise ValueError(f"Unknown model_preference: {model_preference}") print(f"[ORCH_INIT] Creating orchestrator agent with model: {model}") - return Agent( + orchestrator_agent = Agent( model=model, name="Orchestrator Agent", system_prompt=orchestrator_system_prompt, deps_type=orchestrator_deps ) - -@orchestrator_agent.tool -async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Plans the task and assigns it to the appropriate agents""" - try: - logfire.info(f"Planning task: {task}") - # Create a new StreamResponse for Planner Agent - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(planner_stream_output) + @orchestrator_agent.tool + async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Plans the task and assigns it to the appropriate agents""" + try: + logfire.info(f"Planning task: {task}") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Update planner stream - planner_stream_output.steps.append("Planning task...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Run planner agent - agent = planner_agent(model_preference=ctx.deps.model_preference) - planner_response = await agent.run(user_prompt=task) + # Create a new StreamResponse for Planner Agent + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(planner_stream_output) + + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Update planner stream + planner_stream_output.steps.append("Planning task...") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Run planner agent + agent = planner_agent(model_preference=ctx.deps.model_preference) + planner_response = await agent.run(user_prompt=task) - logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") - - # Update planner stream with results - plan_text = planner_response.data.plan - planner_stream_output.steps.append("Task planned successfully") - planner_stream_output.output = plan_text - planner_stream_output.status_code = 200 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Also update orchestrator stream - ctx.deps.stream_output.steps.append("Task planned successfully") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - - return f"Task planned successfully\nTask: {plan_text}" - except Exception as e: - error_msg = f"Error planning task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update planner stream with error - if planner_stream_output: - planner_stream_output.steps.append(f"Planning failed: {str(e)}") - planner_stream_output.status_code = 500 + logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") + + # Update planner stream with results + plan_text = planner_response.data.plan + planner_stream_output.steps.append("Task planned successfully") + planner_stream_output.output = plan_text + planner_stream_output.status_code = 200 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # Also update orchestrator stream - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Planning failed: {str(e)}") + # Also update orchestrator stream + ctx.deps.stream_output.steps.append("Task planned successfully") await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - return f"Failed to plan task: {error_msg}" + return f"Task planned successfully\nTask: {plan_text}" + except Exception as e: + error_msg = f"Error planning task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update planner stream with error + if planner_stream_output: + planner_stream_output.steps.append(f"Planning failed: {str(e)}") + planner_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Also update orchestrator stream + if ctx.deps.stream_output: + ctx.deps.stream_output.steps.append(f"Planning failed: {str(e)}") + await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + + return f"Failed to plan task: {error_msg}" -@orchestrator_agent.tool -async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Assigns coding tasks to the coder agent""" - try: - logfire.info(f"Assigning coding task: {task}") - - # Create a new StreamResponse for Coder Agent - coder_stream_output = StreamResponse( - agent_name="Coder Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(coder_stream_output) - - # Send initial update for Coder Agent - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - # Create deps with the new stream_output - deps_for_coder_agent = CoderAgentDeps( - websocket=ctx.deps.websocket, - stream_output=coder_stream_output, - model_preference = ctx.deps.model_preference - ) - - # Run coder agent - agent = coder_agent(model_preference=ctx.deps.model_preference) - coder_response = await agent.run( - user_prompt=task, - deps=deps_for_coder_agent - ) - logfire.info(f"coder_response: {coder_response}") - logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") - - # Extract response data - response_data = coder_response.data.content - - # Update coder_stream_output with coding results - coder_stream_output.output = response_data - coder_stream_output.status_code = 200 - coder_stream_output.steps.append("Coding task completed successfully") - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - # Add a reminder in the result message to update the plan using planner_agent_update - response_with_reminder = f"{response_data}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" - - return response_with_reminder - except Exception as e: - error_msg = f"Error assigning coding task: {str(e)}" - logfire.error(error_msg, exc_info=True) + @orchestrator_agent.tool + async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Assigns coding tasks to the coder agent""" + try: + logfire.info(f"Assigning coding task: {task}") + + # Create a new StreamResponse for Coder Agent + coder_stream_output = StreamResponse( + agent_name="Coder Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(coder_stream_output) + + # Send initial update for Coder Agent + await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + + # Create deps with the new stream_output + deps_for_coder_agent = CoderAgentDeps( + websocket=ctx.deps.websocket, + stream_output=coder_stream_output, + model_preference = ctx.deps.model_preference + ) + + # Run coder agent + agent = coder_agent(model_preference=ctx.deps.model_preference) + coder_response = await agent.run( + user_prompt=task, + deps=deps_for_coder_agent + ) + logfire.info(f"coder_response: {coder_response}") + logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") + + # Extract response data + response_data = coder_response.data.content + + # Update coder_stream_output with coding results + coder_stream_output.output = response_data + coder_stream_output.status_code = 200 + coder_stream_output.steps.append("Coding task completed successfully") + await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + + # Add a reminder in the result message to update the plan using planner_agent_update + response_with_reminder = f"{response_data}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" + + return response_with_reminder + except Exception as e: + error_msg = f"Error assigning coding task: {str(e)}" + logfire.error(error_msg, exc_info=True) - # Update coder_stream_output with error - coder_stream_output.steps.append(f"Coding task failed: {str(e)}") - coder_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + # Update coder_stream_output with error + coder_stream_output.steps.append(f"Coding task failed: {str(e)}") + coder_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - return f"Failed to assign coding task: {error_msg}" + return f"Failed to assign coding task: {error_msg}" -@orchestrator_agent.tool -async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Assigns web surfing tasks to the web surfer agent""" - try: - logfire.info(f"Assigning web surfing task: {task}") - - # Create a new StreamResponse for WebSurfer - web_surfer_stream_output = StreamResponse( - agent_name="Web Surfer", - instructions=task, - steps=[], - output="", - status_code=0, - live_url=None - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(web_surfer_stream_output) - - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - # Initialize WebSurfer agent with model preference - web_surfer_agent = WebSurfer( - api_url="http://localhost:8000/api/v1/web/stream", - model_preference=ctx.deps.model_preference - ) - logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") - - # Run WebSurfer with its own stream_output - success, message, messages = await web_surfer_agent.generate_reply( - instruction=task, - websocket=ctx.deps.websocket, - stream_output=web_surfer_stream_output - ) - - # Update WebSurfer's stream_output with final result - if success: - web_surfer_stream_output.steps.append("Web search completed successfully") - web_surfer_stream_output.output = message - web_surfer_stream_output.status_code = 200 - - # Add a reminder to update the plan - message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" - else: - web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") + @orchestrator_agent.tool + async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Assigns web surfing tasks to the web surfer agent""" + try: + logfire.info(f"Assigning web surfing task: {task}") + + # Create a new StreamResponse for WebSurfer + web_surfer_stream_output = StreamResponse( + agent_name="Web Surfer", + instructions=task, + steps=[], + output="", + status_code=0, + live_url=None + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(web_surfer_stream_output) + + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + + # Initialize WebSurfer agent with model preference + web_surfer_agent = WebSurfer( + api_url="http://localhost:8000/api/v1/web/stream", + model_preference=ctx.deps.model_preference + ) + logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") + + # Run WebSurfer with its own stream_output + success, message, messages = await web_surfer_agent.generate_reply( + instruction=task, + websocket=ctx.deps.websocket, + stream_output=web_surfer_stream_output + ) + + # Update WebSurfer's stream_output with final result + if success: + web_surfer_stream_output.steps.append("Web search completed successfully") + web_surfer_stream_output.output = message + web_surfer_stream_output.status_code = 200 + + # Add a reminder to update the plan + message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" + else: + web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") + web_surfer_stream_output.status_code = 500 + message_with_reminder = message + + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + + web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + + return message_with_reminder + except Exception as e: + error_msg = f"Error assigning web surfing task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update WebSurfer's stream_output with error + web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") web_surfer_stream_output.status_code = 500 - message_with_reminder = message - - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - return message_with_reminder - except Exception as e: - error_msg = f"Error assigning web surfing task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update WebSurfer's stream_output with error - web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") - web_surfer_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - return f"Failed to assign web surfing task: {error_msg}" - -@orchestrator_agent.tool -async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: - """Sends a question to the frontend and waits for human input""" - try: - logfire.info(f"Asking human: {question}") - - # Create a new StreamResponse for Human Input - human_stream_output = StreamResponse( - agent_name="Human Input", - instructions=question, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(human_stream_output) - - # Send the question to frontend - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - # Update stream with waiting message - human_stream_output.steps.append("Waiting for human input...") - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - # Wait for response from frontend - response = await ctx.deps.websocket.receive_text() - - # Update stream with response - human_stream_output.steps.append("Received human input") - human_stream_output.output = response - human_stream_output.status_code = 200 - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - return response - except Exception as e: - error_msg = f"Error getting human input: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update stream with error - human_stream_output.steps.append(f"Failed to get human input: {str(e)}") - human_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - return f"Failed to get human input: {error_msg}" + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + return f"Failed to assign web surfing task: {error_msg}" -@orchestrator_agent.tool -async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_task: str) -> str: - """ - Updates the todo.md file to mark a task as completed and returns the full updated plan. - - Args: - completed_task: Description of the completed task including which agent performed it - - Returns: - The complete updated todo.md content with tasks marked as completed - """ - try: - logfire.info(f"Updating plan with completed task: {completed_task}") - - # Create a new StreamResponse for Planner Agent update - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=f"Update todo.md to mark as completed: {completed_task}", - steps=[], - output="", - status_code=0 - ) - - # Send initial update - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Directly read and update the todo.md file - base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) - planner_dir = os.path.join(base_dir, "agents", "planner") - todo_path = os.path.join(planner_dir, "todo.md") - - planner_stream_output.steps.append("Reading current todo.md...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Make sure the directory exists - os.makedirs(planner_dir, exist_ok=True) - + @orchestrator_agent.tool + async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: + """Sends a question to the frontend and waits for human input""" try: - # Check if todo.md exists - if not os.path.exists(todo_path): - planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # We'll directly call planner_agent.run() to create a new plan first - plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - agent = planner_agent(model_preference=ctx.deps.model_preference) - plan_response = await agent.run(user_prompt=plan_prompt) - current_content = plan_response.data.plan - else: - # Read existing todo.md - with open(todo_path, "r") as file: - current_content = file.read() - planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + logfire.info(f"Asking human: {question}") - # Now call planner_agent.run() with specific instructions to update the plan - update_prompt = f""" - Here is the current todo.md content: + # Create a new StreamResponse for Human Input + human_stream_output = StreamResponse( + agent_name="Human Input", + instructions=question, + steps=[], + output="", + status_code=0 + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(human_stream_output) + + # Send the question to frontend + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - {current_content} + # Update stream with waiting message + human_stream_output.steps.append("Waiting for human input...") + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - Please update this plan to mark the following task as completed: {completed_task} - Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. - """ + # Wait for response from frontend + response = await ctx.deps.websocket.receive_text() - planner_stream_output.steps.append("Asking planner to update the plan...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + # Update stream with response + human_stream_output.steps.append("Received human input") + human_stream_output.output = response + human_stream_output.status_code = 200 + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - updated_plan_response = await planner_agent.run(user_prompt=update_prompt) - updated_plan = updated_plan_response.data.plan + return response + except Exception as e: + error_msg = f"Error getting human input: {str(e)}" + logfire.error(error_msg, exc_info=True) - # Write the updated plan back to todo.md - with open(todo_path, "w") as file: - file.write(updated_plan) + # Update stream with error + human_stream_output.steps.append(f"Failed to get human input: {str(e)}") + human_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - planner_stream_output.steps.append("Plan updated successfully") - planner_stream_output.output = updated_plan - planner_stream_output.status_code = 200 + return f"Failed to get human input: {error_msg}" + + @orchestrator_agent.tool + async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_task: str) -> str: + """ + Updates the todo.md file to mark a task as completed and returns the full updated plan. + + Args: + completed_task: Description of the completed task including which agent performed it + + Returns: + The complete updated todo.md content with tasks marked as completed + """ + try: + logfire.info(f"Updating plan with completed task: {completed_task}") + + # Create a new StreamResponse for Planner Agent update + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=f"Update todo.md to mark as completed: {completed_task}", + steps=[], + output="", + status_code=0 + ) + + # Send initial update await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # Update orchestrator stream - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Plan updated to mark task as completed: {completed_task}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + # Directly read and update the todo.md file + base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) + planner_dir = os.path.join(base_dir, "agents", "planner") + todo_path = os.path.join(planner_dir, "todo.md") + + planner_stream_output.steps.append("Reading current todo.md...") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Make sure the directory exists + os.makedirs(planner_dir, exist_ok=True) - return updated_plan + try: + # Check if todo.md exists + if not os.path.exists(todo_path): + planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # We'll directly call planner_agent.run() to create a new plan first + plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" + agent = planner_agent(model_preference=ctx.deps.model_preference) + plan_response = await agent.run(user_prompt=plan_prompt) + current_content = plan_response.data.plan + + logfire.info(f"Updating plan task using model type: {ctx.deps.model_preference}") + + else: + # Read existing todo.md + with open(todo_path, "r") as file: + current_content = file.read() + planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Now call planner_agent.run() with specific instructions to update the plan + update_prompt = f""" + Here is the current todo.md content: + + {current_content} + + Please update this plan to mark the following task as completed: {completed_task} + Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. + """ + + planner_stream_output.steps.append("Asking planner to update the plan...") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + agent = planner_agent(model_preference=ctx.deps.model_preference) + updated_plan_response = await agent.run(user_prompt=update_prompt) + updated_plan = updated_plan_response.data.plan + + logfire.info(f"Updating prompt for plan task using model type: {ctx.deps.model_preference}") + + # Write the updated plan back to todo.md + with open(todo_path, "w") as file: + file.write(updated_plan) + + planner_stream_output.steps.append("Plan updated successfully") + planner_stream_output.output = updated_plan + planner_stream_output.status_code = 200 + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Update orchestrator stream + if ctx.deps.stream_output: + ctx.deps.stream_output.steps.append(f"Plan updated to mark task as completed: {completed_task}") + await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + + return updated_plan + + except Exception as e: + error_msg = f"Error during plan update operations: {str(e)}" + logfire.error(error_msg, exc_info=True) + + planner_stream_output.steps.append(f"Plan update failed: {str(e)}") + planner_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + return f"Failed to update the plan: {error_msg}" except Exception as e: - error_msg = f"Error during plan update operations: {str(e)}" + error_msg = f"Error updating plan: {str(e)}" logfire.error(error_msg, exc_info=True) - planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + # Update stream output with error + if ctx.deps.stream_output: + ctx.deps.stream_output.steps.append(f"Failed to update plan: {str(e)}") + await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - return f"Failed to update the plan: {error_msg}" - - except Exception as e: - error_msg = f"Error updating plan: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update stream output with error - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Failed to update plan: {str(e)}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - - return f"Failed to update plan: {error_msg}" + return f"Failed to update plan: {error_msg}" + + logfire.info("All tools initialized for orchestrator agent") + return orchestrator_agent + # Helper function for sending WebSocket messages async def _safe_websocket_send(websocket: Optional[WebSocket], message: Any) -> bool: diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 53592ae..5d29b7d 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -179,7 +179,7 @@ class PlannerResult(BaseModel): -def planner_agent(model_preference: str = "Anthropic"): +async def planner_agent(model_preference: str = "Anthropic") -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -193,162 +193,164 @@ def planner_agent(model_preference: str = "Anthropic"): result_type=PlannerResult, system_prompt=planner_prompt ) - return planner_agent -@planner_agent.tool_plain -async def update_todo_status(task_description: str) -> str: - """ - A helper function that logs the update request but lets the planner agent handle the actual update logic. - - Args: - task_description: Description of the completed task + @planner_agent.tool_plain + async def update_todo_status(task_description: str) -> str: + """ + A helper function that logs the update request but lets the planner agent handle the actual update logic. - Returns: - A simple log message - """ - logfire.info(f"Received request to update todo.md for task: {task_description}") - return f"Received update request for: {task_description}" + Args: + task_description: Description of the completed task + + Returns: + A simple log message + """ + logfire.info(f"Received request to update todo.md for task: {task_description}") + return f"Received update request for: {task_description}" -@planner_agent.tool_plain -async def execute_terminal(command: str) -> str: - """ - Executes a terminal command within the planner directory for file operations. - This consolidated tool handles reading and writing plan files. - Restricted to only read and write operations for security. - """ - try: - logfire.info(f"Executing terminal command: {command}") - - # Define the restricted directory - base_dir = os.path.abspath(os.path.dirname(__file__)) - planner_dir = os.path.join(base_dir, "planner") - os.makedirs(planner_dir, exist_ok=True) - - # Extract the base command - base_command = command.split()[0] - - # Allow only read and write operations - ALLOWED_COMMANDS = {"cat", "echo", "ls"} - - # Security checks - if base_command not in ALLOWED_COMMANDS: - return f"Error: Command '{base_command}' is not allowed. Only read and write operations are permitted." - - if ".." in command or "~" in command or "/" in command: - return "Error: Path traversal attempts are not allowed." - - # Change to the restricted directory - original_dir = os.getcwd() - os.chdir(planner_dir) - + @planner_agent.tool_plain + async def execute_terminal(command: str) -> str: + """ + Executes a terminal command within the planner directory for file operations. + This consolidated tool handles reading and writing plan files. + Restricted to only read and write operations for security. + """ try: - # Handle echo with >> (append) - if base_command == "echo" and ">>" in command: - try: - # Split only on the first occurrence of >> - parts = command.split(">>", 1) - echo_part = parts[0].strip() + logfire.info(f"Executing terminal command: {command}") + + # Define the restricted directory + base_dir = os.path.abspath(os.path.dirname(__file__)) + planner_dir = os.path.join(base_dir, "planner") + os.makedirs(planner_dir, exist_ok=True) + + # Extract the base command + base_command = command.split()[0] + + # Allow only read and write operations + ALLOWED_COMMANDS = {"cat", "echo", "ls"} + + # Security checks + if base_command not in ALLOWED_COMMANDS: + return f"Error: Command '{base_command}' is not allowed. Only read and write operations are permitted." + + if ".." in command or "~" in command or "/" in command: + return "Error: Path traversal attempts are not allowed." + + # Change to the restricted directory + original_dir = os.getcwd() + os.chdir(planner_dir) + + try: + # Handle echo with >> (append) + if base_command == "echo" and ">>" in command: + try: + # Split only on the first occurrence of >> + parts = command.split(">>", 1) + echo_part = parts[0].strip() + file_path = parts[1].strip() + + # Extract content after echo command + content = echo_part[4:].strip() + + # Handle quotes if present + if (content.startswith('"') and content.endswith('"')) or \ + (content.startswith("'") and content.endswith("'")): + content = content[1:-1] + + # Append to file + with open(file_path, "a") as file: + file.write(content + "\n") + return f"Successfully appended to {file_path}" + except Exception as e: + logfire.error(f"Error appending to file: {str(e)}", exc_info=True) + return f"Error appending to file: {str(e)}" + + # Special handling for echo with redirection (file writing) + elif ">" in command and base_command == "echo" and ">>" not in command: + # Simple parsing for echo "content" > file.txt + parts = command.split(">", 1) + echo_cmd = parts[0].strip() file_path = parts[1].strip() - # Extract content after echo command - content = echo_part[4:].strip() - - # Handle quotes if present + # Extract content between echo and > (removing quotes if present) + content = echo_cmd[5:].strip() if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): + (content.startswith("'") and content.endswith("'")): content = content[1:-1] - # Append to file - with open(file_path, "a") as file: - file.write(content + "\n") - return f"Successfully appended to {file_path}" - except Exception as e: - logfire.error(f"Error appending to file: {str(e)}", exc_info=True) - return f"Error appending to file: {str(e)}" - - # Special handling for echo with redirection (file writing) - elif ">" in command and base_command == "echo" and ">>" not in command: - # Simple parsing for echo "content" > file.txt - parts = command.split(">", 1) - echo_cmd = parts[0].strip() - file_path = parts[1].strip() + # Write to file + try: + with open(file_path, "w") as file: + file.write(content) + return f"Successfully wrote to {file_path}" + except Exception as e: + logfire.error(f"Error writing to file: {str(e)}", exc_info=True) + return f"Error writing to file: {str(e)}" - # Extract content between echo and > (removing quotes if present) - content = echo_cmd[5:].strip() - if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): - content = content[1:-1] - - # Write to file - try: - with open(file_path, "w") as file: - file.write(content) - return f"Successfully wrote to {file_path}" - except Exception as e: - logfire.error(f"Error writing to file: {str(e)}", exc_info=True) - return f"Error writing to file: {str(e)}" - - # Handle cat with here-document for multiline file writing - elif "<<" in command and base_command == "cat": - try: - # Parse the command: cat > file.md << 'EOF'\nplan content\nEOF - cmd_parts = command.split("<<", 1) - cat_part = cmd_parts[0].strip() - doc_part = cmd_parts[1].strip() - - # Extract filename - if ">" in cat_part: - file_path = cat_part.split(">", 1)[1].strip() - else: - return "Error: Invalid cat command format. Must include redirection." - - # Parse the heredoc content - if "\n" in doc_part: - delimiter_and_content = doc_part.split("\n", 1) - delimiter = delimiter_and_content[0].strip("'").strip('"') - content = delimiter_and_content[1] + # Handle cat with here-document for multiline file writing + elif "<<" in command and base_command == "cat": + try: + # Parse the command: cat > file.md << 'EOF'\nplan content\nEOF + cmd_parts = command.split("<<", 1) + cat_part = cmd_parts[0].strip() + doc_part = cmd_parts[1].strip() + + # Extract filename + if ">" in cat_part: + file_path = cat_part.split(">", 1)[1].strip() + else: + return "Error: Invalid cat command format. Must include redirection." - # Find the end delimiter and extract content - if f"\n{delimiter}" in content: - content = content.split(f"\n{delimiter}")[0] + # Parse the heredoc content + if "\n" in doc_part: + delimiter_and_content = doc_part.split("\n", 1) + delimiter = delimiter_and_content[0].strip("'").strip('"') + content = delimiter_and_content[1] - # Write to file - with open(file_path, "w") as file: - file.write(content) - return f"Successfully wrote multiline content to {file_path}" + # Find the end delimiter and extract content + if f"\n{delimiter}" in content: + content = content.split(f"\n{delimiter}")[0] + + # Write to file + with open(file_path, "w") as file: + file.write(content) + return f"Successfully wrote multiline content to {file_path}" + else: + return "Error: End delimiter not found in heredoc" else: - return "Error: End delimiter not found in heredoc" - else: - return "Error: Invalid heredoc format" - except Exception as e: - logfire.error(f"Error processing cat with heredoc: {str(e)}", exc_info=True) - return f"Error processing cat with heredoc: {str(e)}" - - # Handle cat for reading files - elif base_command == "cat" and ">" not in command and "<<" not in command: - try: - file_path = command.split()[1] - with open(file_path, "r") as file: - content = file.read() - return content - except Exception as e: - logfire.error(f"Error reading file: {str(e)}", exc_info=True) - return f"Error reading file: {str(e)}" - - # Handle ls for listing files - elif base_command == "ls": - try: - files = os.listdir('.') - return "Files in planner directory:\n" + "\n".join(files) - except Exception as e: - logfire.error(f"Error listing files: {str(e)}", exc_info=True) - return f"Error listing files: {str(e)}" - else: - return f"Error: Command '{command}' is not supported. Only read and write operations are permitted." - - finally: - os.chdir(original_dir) - - except Exception as e: - logfire.error(f"Error executing command: {str(e)}", exc_info=True) - return f"Error executing command: {str(e)}" \ No newline at end of file + return "Error: Invalid heredoc format" + except Exception as e: + logfire.error(f"Error processing cat with heredoc: {str(e)}", exc_info=True) + return f"Error processing cat with heredoc: {str(e)}" + + # Handle cat for reading files + elif base_command == "cat" and ">" not in command and "<<" not in command: + try: + file_path = command.split()[1] + with open(file_path, "r") as file: + content = file.read() + return content + except Exception as e: + logfire.error(f"Error reading file: {str(e)}", exc_info=True) + return f"Error reading file: {str(e)}" + + # Handle ls for listing files + elif base_command == "ls": + try: + files = os.listdir('.') + return "Files in planner directory:\n" + "\n".join(files) + except Exception as e: + logfire.error(f"Error listing files: {str(e)}", exc_info=True) + return f"Error listing files: {str(e)}" + else: + return f"Error: Command '{command}' is not supported. Only read and write operations are permitted." + + finally: + os.chdir(original_dir) + + except Exception as e: + logfire.error(f"Error executing command: {str(e)}", exc_info=True) + return f"Error executing command: {str(e)}" + + logfire.info("All tools initialized for planner agent") + return planner_agent \ No newline at end of file diff --git a/cortex_on/main.py b/cortex_on/main.py index 4bab62f..4488c3a 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -29,9 +29,9 @@ async def lifespan(app: FastAPI): yield - # Cleanup - if instructor: - await instructor.shutdown() + # # Cleanup + # if instructor: + # await instructor.shutdown() app: FastAPI = FastAPI(lifespan=lifespan) From 708e833c820fcf87249a595d0b01a6f1c3d4a5d6 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Wed, 16 Apr 2025 22:59:52 +0530 Subject: [PATCH 05/24] feat: agent reinitialization everytime the model_preference changes --- cortex_on/agents/hello.py | 65 -------------------------- cortex_on/agents/orchestrator_agent.py | 16 ++++--- cortex_on/instructor.py | 4 +- ta-browser/core/server/routes/web.py | 9 ---- 4 files changed, 11 insertions(+), 83 deletions(-) delete mode 100644 cortex_on/agents/hello.py diff --git a/cortex_on/agents/hello.py b/cortex_on/agents/hello.py deleted file mode 100644 index 2c453cc..0000000 --- a/cortex_on/agents/hello.py +++ /dev/null @@ -1,65 +0,0 @@ - -@dataclass -class orchestrator_deps: - websocket: Optional[WebSocket] = None - stream_output: Optional[StreamResponse] = None - # Add a collection to track agent-specific streams - agent_responses: Optional[List[StreamResponse]] = None - model_preference: str = "Anthropic" - -async def get_model_preference() -> str: - """Fetch model preference from ta_browser API""" - try: - async with aiohttp.ClientSession() as session: - async with session.get("http://localhost:8000/api/v1/model/preference") as response: - if response.status == 200: - data = await response.json() - return data.get("model_preference", "Anthropic") - else: - logfire.error(f"Failed to get model preference: {response.status}") - return "Anthropic" - except Exception as e: - logfire.error(f"Error getting model preference: {str(e)}") - return "Anthropic" - -async def initialize_orchestrator_agent(): - """Initialize the orchestrator agent with model preference from ta_browser""" - model_preference = await get_model_preference() - model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() - ) - - print(f"Orchestrator agent model initialized with MODEL PREFERENCE: {model_preference}") - - return Agent( - model=model, - name="Orchestrator Agent", - system_prompt=orchestrator_system_prompt, - deps_type=orchestrator_deps - ) - -# Initialize the agent -orchestrator_agent = None - -def get_orchestrator_agent(): - """Get the orchestrator agent, initializing it if necessary""" - global orchestrator_agent - if orchestrator_agent is None: - raise RuntimeError("Orchestrator agent not initialized. Call initialize_orchestrator_agent() first.") - return orchestrator_agent - -@get_orchestrator_agent().tool -async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Plans the task and assigns it to the appropriate agents""" - try: - logfire.info(f"Planning task: {task}") - - # Create a new StreamResponse for Planner Agent - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 4a5eec5..2bfda6a 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -12,7 +12,7 @@ from pydantic_ai import Agent, RunContext from agents.web_surfer import WebSurfer from utils.stream_response_format import StreamResponse -from agents.planner_agent import planner_agent, update_todo_status +from agents.planner_agent import planner_agent from agents.code_agent import coder_agent, CoderAgentDeps from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client @@ -181,7 +181,7 @@ class orchestrator_deps: # deps_type=orchestrator_deps # ) # print("[ORCH_INIT] Orchestrator agent created successfully") -async def orchestrator_agent(model_preference: str) -> Agent: +async def orchestrator_agent(model_preference: str) : if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -222,7 +222,7 @@ async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) # Run planner agent - agent = planner_agent(model_preference=ctx.deps.model_preference) + agent = await planner_agent(model_preference=ctx.deps.model_preference) planner_response = await agent.run(user_prompt=task) logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") @@ -286,7 +286,7 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: ) # Run coder agent - agent = coder_agent(model_preference=ctx.deps.model_preference) + agent = await coder_agent(model_preference=ctx.deps.model_preference) coder_response = await agent.run( user_prompt=task, deps=deps_for_coder_agent @@ -475,7 +475,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas # We'll directly call planner_agent.run() to create a new plan first plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - agent = planner_agent(model_preference=ctx.deps.model_preference) + + agent = await planner_agent(model_preference=ctx.deps.model_preference) plan_response = await agent.run(user_prompt=plan_prompt) current_content = plan_response.data.plan @@ -501,7 +502,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas planner_stream_output.steps.append("Asking planner to update the plan...") await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - agent = planner_agent(model_preference=ctx.deps.model_preference) + agent = await planner_agent(model_preference=ctx.deps.model_preference) updated_plan_response = await agent.run(user_prompt=update_prompt) updated_plan = updated_plan_response.data.plan @@ -544,7 +545,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas return f"Failed to update plan: {error_msg}" - logfire.info("All tools initialized for orchestrator agent") + logfire.info("All tools initialized for orchestrator agent") + print(f"[ORCH_INIT] Orchestrator agent initialized successfully: {orchestrator_agent}") return orchestrator_agent diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index d92cf4e..4821f80 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -28,7 +28,7 @@ print("[INIT] Loading agents...") print(f"[INIT] Planner agent loaded: {planner_agent}") print(f"[INIT] Coder agent loaded: {coder_agent}") -print(f"[INIT] Orchestrator agent loaded: {get_orchestrator_agent}") +print(f"[INIT] Orchestrator agent loaded: {orchestrator_agent}") class DateTimeEncoder(json.JSONEncoder): """Custom JSON encoder that can handle datetime objects""" @@ -97,7 +97,7 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: print("Orchestrator agent running") - agent = orchestrator_agent(self.model_preference) + agent = await orchestrator_agent(self.model_preference) orchestrator_response = await agent.run( user_prompt=task, deps=deps_for_orchestrator diff --git a/ta-browser/core/server/routes/web.py b/ta-browser/core/server/routes/web.py index 8dabe34..1be01af 100644 --- a/ta-browser/core/server/routes/web.py +++ b/ta-browser/core/server/routes/web.py @@ -154,12 +154,3 @@ async def stream_session( logger.error(f"Error in /stream: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) -@router.get("/api/v1/model/preference") -async def get_model_preference(): - """Get the current model preference""" - try: - # Get model preference from environment or use default - model_preference = os.getenv("MODEL_PREFERENCE", "Anthropic") - return {"model_preference": model_preference} - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) \ No newline at end of file From 79edd97a06ff7734f90ee12cd90c69c393d80696 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Fri, 18 Apr 2025 15:36:46 +0530 Subject: [PATCH 06/24] feat: fixed websocket error --- cortex_on/instructor.py | 12 +++++++----- cortex_on/main.py | 17 +++++++++++++---- frontend/src/components/home/ChatList.tsx | 22 +++++++++++++++++++++- frontend/src/components/home/Header.tsx | 2 +- 4 files changed, 42 insertions(+), 11 deletions(-) diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index 4821f80..940d11c 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -128,14 +128,16 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: finally: logfire.info("Orchestration process complete") - # Clear any sensitive data + + async def shutdown(self): """Clean shutdown of orchestrator""" try: - # Close websocket if open - if self.websocket: - await self.websocket.close() - + #Close websocket if open + # if self.websocket: + # await self.websocket.close() + # logfire.info("websocket shutdown complete") + # print("websocket shutdown complete") # Clear all responses self.orchestrator_response = [] diff --git a/cortex_on/main.py b/cortex_on/main.py index 4488c3a..55926b1 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -95,7 +95,16 @@ async def websocket_endpoint(websocket: WebSocket): await websocket.accept() model_preference = app.state.model_preference print(f"[WEBSOCKET] New connection using model preference: {model_preference}") - while True: - data = await websocket.receive_text() - print(f"[WEBSOCKET] Received message, using model: {model_preference}") - await generate_response(data, websocket, model_preference) + try: + while True: + try: + data = await websocket.receive_text() + print(f"[WEBSOCKET] Received message, using model: {model_preference}") + await generate_response(data, websocket, model_preference) + except Exception as e: + print(f"[WEBSOCKET] Error handling message: {str(e)}") + break + except Exception as e: + print(f"[WEBSOCKET] Connection error: {str(e)}") + finally: + print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/ChatList.tsx b/frontend/src/components/home/ChatList.tsx index 9d8cb21..5c6bc7c 100644 --- a/frontend/src/components/home/ChatList.tsx +++ b/frontend/src/components/home/ChatList.tsx @@ -94,9 +94,29 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { retryOnError: true, shouldReconnect: () => true, reconnectInterval: 3000, + share: true, } ); + useEffect(() => { + if (readyState === ReadyState.CONNECTING) { + if (messages.length > 0) { + setIsLoading(true); + } + } else if (readyState === ReadyState.OPEN) { + setIsLoading(false); + } + }, [readyState, messages.length]); + + useEffect(() => { + if (messages.length === 0) { + setLiveUrl(""); + setCurrentOutput(null); + setOutputsList([]); + setIsLoading(false); + } + }, [messages.length]); + const scrollToBottom = (smooth = true) => { if (scrollAreaRef.current) { const scrollableDiv = scrollAreaRef.current.querySelector( @@ -557,7 +577,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { className="space-y-2 animate-fade-in animate-once animate-duration-500" key={idx} > - {readyState === ReadyState.CONNECTING ? ( + {readyState === ReadyState.CONNECTING && messages.length > 0 ? ( <> diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 5fba65a..d91c194 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - nav("/"); + window.location.href = "/"; }} > From 176d71ab6eb0c1f7d6cbc36bded0fa06f0b68510 Mon Sep 17 00:00:00 2001 From: aryan Date: Fri, 18 Apr 2025 17:12:39 +0530 Subject: [PATCH 07/24] fix(logging): Replace print statements with logfire for improved logging consistency across the application --- cortex_on/agents/orchestrator_agent.py | 2 +- cortex_on/agents/web_surfer.py | 2 +- cortex_on/instructor.py | 2 +- cortex_on/main.py | 31 ++++++++++++++----------- frontend/src/components/home/Header.tsx | 2 +- 5 files changed, 22 insertions(+), 17 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 2bfda6a..bdb177a 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -181,7 +181,7 @@ class orchestrator_deps: # deps_type=orchestrator_deps # ) # print("[ORCH_INIT] Orchestrator agent created successfully") -async def orchestrator_agent(model_preference: str) : +async def orchestrator_agent(model_preference: str) -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 03503d5..2711348 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -47,7 +47,7 @@ async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, An "critique_disabled": False, "model_preference": self.model_preference } - print(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") + logfire.info(f"Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index 940d11c..388a7bb 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -41,7 +41,7 @@ def default(self, obj): # Main Orchestrator Class class SystemInstructor: def __init__(self, model_preference: str = "Anthropic"): - print(f"[INIT] Initializing SystemInstructor with model_preference: {model_preference}") + logfire.info(f"Initializing SystemInstructor with model_preference: {model_preference}") self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None self.orchestrator_response: List[StreamResponse] = [] diff --git a/cortex_on/main.py b/cortex_on/main.py index 55926b1..deb1f38 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -3,8 +3,9 @@ from contextlib import asynccontextmanager # Third-party imports -from fastapi import FastAPI, WebSocket, Depends +from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends from fastapi.middleware.cors import CORSMiddleware +import logfire # Local application imports from instructor import SystemInstructor @@ -20,7 +21,7 @@ async def lifespan(app: FastAPI): # Set default model preference at startup app.state.model_preference = MODEL_PREFERENCE - print(f"[STARTUP] Setting default model preference to: {MODEL_PREFERENCE}") + logfire.info(f"Setting default model preference to: {MODEL_PREFERENCE}") # Initialize the instructor global instructor @@ -49,7 +50,7 @@ async def get_model_preference() -> str: """ Get the current model preference from app state """ - print(f"[GET] Current model preference: {app.state.model_preference}") + logfire.info(f"Current model preference: {app.state.model_preference}") return app.state.model_preference @app.get("/set_model_preference") @@ -58,10 +59,10 @@ async def set_model_preference(model: str): Set the model preference (Anthropic or OpenAI) and reinitialize the instructor """ if model not in ["Anthropic", "OpenAI"]: - print(f"[ERROR] Invalid model preference attempted: {model}") + logfire.error(f"Invalid model preference attempted: {model}") return {"error": "Model must be 'Anthropic' or 'OpenAI'"} - print(f"[SET] Changing model preference from {app.state.model_preference} to {model}") + logfire.info(f"Changing model preference from {app.state.model_preference} to {model}") app.state.model_preference = model # Reinitialize the instructor with new model preference @@ -69,14 +70,14 @@ async def set_model_preference(model: str): if instructor: await instructor.shutdown() instructor = SystemInstructor(model_preference=model) - print(f"[SET] Instructor reinitialized with model preference: {model}") + logfire.info(f"Instructor reinitialized with model preference: {model}") return {"message": f"Model preference set to {model}"} async def generate_response(task: str, websocket: Optional[WebSocket] = None, model_preference: str = None): if model_preference is None: model_preference = app.state.model_preference - print(f"[GENERATE] Using model preference: {model_preference} for task: {task[:30]}...") + logfire.info(f"Using model preference: {model_preference} for task: {task[:30]}...") global instructor if not instructor: @@ -86,7 +87,7 @@ async def generate_response(task: str, websocket: Optional[WebSocket] = None, mo @app.get("/agent/chat") async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List: - print(f"[CHAT] Received chat request with model preference: {model_preference}") + logfire.info(f"Received chat request with model preference: {model_preference}") final_agent_response = await generate_response(task, model_preference=model_preference) return final_agent_response @@ -94,17 +95,21 @@ async def agent_chat(task: str, model_preference: str = Depends(get_model_prefer async def websocket_endpoint(websocket: WebSocket): await websocket.accept() model_preference = app.state.model_preference - print(f"[WEBSOCKET] New connection using model preference: {model_preference}") + logfire.info(f"New connection using model preference: {model_preference}") try: while True: try: + # Use a timeout to detect if the client is still connected data = await websocket.receive_text() - print(f"[WEBSOCKET] Received message, using model: {model_preference}") + logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) + except WebSocketDisconnect: + print("[WEBSOCKET] Client disconnected") + break except Exception as e: print(f"[WEBSOCKET] Error handling message: {str(e)}") - break + if "disconnect message has been received" in str(e): + print("[WEBSOCKET] Disconnect detected, closing connection") + break except Exception as e: print(f"[WEBSOCKET] Connection error: {str(e)}") - finally: - print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index d91c194..5fba65a 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - window.location.href = "/"; + nav("/"); }} > From a88b197cc0eafa91bce6e77c4bf4a18b4d500028 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Tue, 22 Apr 2025 22:35:44 +0530 Subject: [PATCH 08/24] feat[openai compatibility]:made changes that CortexON can now switches between OpenAI and Antropic models --- cortex_on/agents/orchestrator_agent.py | 49 +++++++--------------- cortex_on/agents/planner_agent.py | 6 +-- cortex_on/agents/web_surfer.py | 2 +- cortex_on/main.py | 13 +++--- frontend/src/components/home/ChatList.tsx | 3 +- frontend/src/components/home/Header.tsx | 2 +- ta-browser/core/utils/openai_msg_parser.py | 1 + 7 files changed, 27 insertions(+), 49 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index bdb177a..c344213 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -16,8 +16,6 @@ from agents.code_agent import coder_agent, CoderAgentDeps from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client -print("[ORCH_INIT] Starting orchestrator agent initialization") - @dataclass class orchestrator_deps: websocket: Optional[WebSocket] = None @@ -26,7 +24,6 @@ class orchestrator_deps: agent_responses: Optional[List[StreamResponse]] = None model_preference: str = "Anthropic" - orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. [AGENT CAPABILITIES] @@ -162,25 +159,6 @@ class orchestrator_deps: - Format: "Task description (agent_name)" """ -# print("[ORCH_INIT] Getting Anthropic client...") -# client = get_client() -# print(f"[ORCH_INIT] Anthropic client obtained: {client}") - -# print("[ORCH_INIT] Initializing Anthropic model...") -# model = get_anthropic_model_instance() - -# print(f"[ORCH_INIT] Anthropic model initialized: {model}") - -# print("Orchestrator agent model initialized without MODEL PREFERENCE") - -# print("[ORCH_INIT] Creating orchestrator agent...") -# orchestrator_agent = Agent( -# model=model, -# name="Orchestrator Agent", -# system_prompt=orchestrator_system_prompt, -# deps_type=orchestrator_deps -# ) -# print("[ORCH_INIT] Orchestrator agent created successfully") async def orchestrator_agent(model_preference: str) -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() @@ -195,7 +173,7 @@ async def orchestrator_agent(model_preference: str) -> Agent: system_prompt=orchestrator_system_prompt, deps_type=orchestrator_deps ) - + @orchestrator_agent.tool async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: """Plans the task and assigns it to the appropriate agents""" @@ -226,7 +204,7 @@ async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: planner_response = await agent.run(user_prompt=task) logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") - + # Update planner stream with results plan_text = planner_response.data.plan planner_stream_output.steps.append("Task planned successfully") @@ -281,8 +259,7 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: # Create deps with the new stream_output deps_for_coder_agent = CoderAgentDeps( websocket=ctx.deps.websocket, - stream_output=coder_stream_output, - model_preference = ctx.deps.model_preference + stream_output=coder_stream_output ) # Run coder agent @@ -294,6 +271,7 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: logfire.info(f"coder_response: {coder_response}") logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") + # Extract response data response_data = coder_response.data.content @@ -346,7 +324,7 @@ async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: model_preference=ctx.deps.model_preference ) logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") - + # Run WebSurfer with its own stream_output success, message, messages = await web_surfer_agent.generate_reply( instruction=task, @@ -479,9 +457,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas agent = await planner_agent(model_preference=ctx.deps.model_preference) plan_response = await agent.run(user_prompt=plan_prompt) current_content = plan_response.data.plan - logfire.info(f"Updating plan task using model type: {ctx.deps.model_preference}") - + else: # Read existing todo.md with open(todo_path, "r") as file: @@ -502,12 +479,12 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas planner_stream_output.steps.append("Asking planner to update the plan...") await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - agent = await planner_agent(model_preference=ctx.deps.model_preference) - updated_plan_response = await agent.run(user_prompt=update_prompt) + updated_plan_response = await planner_agent.run(user_prompt=update_prompt) updated_plan = updated_plan_response.data.plan + logfire.info(f"Updating prompt for plan task using model type: {ctx.deps.model_preference}") - + # Write the updated plan back to todo.md with open(todo_path, "w") as file: file.write(updated_plan) @@ -529,7 +506,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.error(error_msg, exc_info=True) planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = 500 + planner_stream_output.status_code = a500 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) return f"Failed to update the plan: {error_msg}" @@ -544,12 +521,16 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) return f"Failed to update plan: {error_msg}" - + + + logfire.info("All tools initialized for orchestrator agent") print(f"[ORCH_INIT] Orchestrator agent initialized successfully: {orchestrator_agent}") + return orchestrator_agent + # Helper function for sending WebSocket messages async def _safe_websocket_send(websocket: Optional[WebSocket], message: Any) -> bool: """Safely send message through websocket with error handling""" diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 5d29b7d..d5d7832 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -176,9 +176,6 @@ class PlannerResult(BaseModel): plan: str = Field(description="The generated or updated plan in string format - this should be the complete plan text") - - - async def planner_agent(model_preference: str = "Anthropic") -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() @@ -351,6 +348,7 @@ async def execute_terminal(command: str) -> str: except Exception as e: logfire.error(f"Error executing command: {str(e)}", exc_info=True) return f"Error executing command: {str(e)}" - + + logfire.info("All tools initialized for planner agent") return planner_agent \ No newline at end of file diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 2711348..e9fe772 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -47,7 +47,7 @@ async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, An "critique_disabled": False, "model_preference": self.model_preference } - logfire.info(f"Making API call with model_preference: {self.model_preference}") + logfire.info(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/main.py b/cortex_on/main.py index deb1f38..bb3aa93 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -3,14 +3,13 @@ from contextlib import asynccontextmanager # Third-party imports -from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends +from fastapi import FastAPI, WebSocket, WebSocketDisconnect ,Depends from fastapi.middleware.cors import CORSMiddleware import logfire # Local application imports from instructor import SystemInstructor - # Default model preference is Anthropic MODEL_PREFERENCE = "Anthropic" @@ -22,7 +21,7 @@ async def lifespan(app: FastAPI): # Set default model preference at startup app.state.model_preference = MODEL_PREFERENCE logfire.info(f"Setting default model preference to: {MODEL_PREFERENCE}") - + # Initialize the instructor global instructor instructor = SystemInstructor(model_preference=MODEL_PREFERENCE) @@ -99,17 +98,15 @@ async def websocket_endpoint(websocket: WebSocket): try: while True: try: - # Use a timeout to detect if the client is still connected data = await websocket.receive_text() logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) except WebSocketDisconnect: - print("[WEBSOCKET] Client disconnected") + print("Websocket disconnected") break except Exception as e: print(f"[WEBSOCKET] Error handling message: {str(e)}") - if "disconnect message has been received" in str(e): - print("[WEBSOCKET] Disconnect detected, closing connection") - break except Exception as e: print(f"[WEBSOCKET] Connection error: {str(e)}") + # finally: + # print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/ChatList.tsx b/frontend/src/components/home/ChatList.tsx index 5c6bc7c..50fba0a 100644 --- a/frontend/src/components/home/ChatList.tsx +++ b/frontend/src/components/home/ChatList.tsx @@ -117,6 +117,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { } }, [messages.length]); + const scrollToBottom = (smooth = true) => { if (scrollAreaRef.current) { const scrollableDiv = scrollAreaRef.current.querySelector( @@ -577,7 +578,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { className="space-y-2 animate-fade-in animate-once animate-duration-500" key={idx} > - {readyState === ReadyState.CONNECTING && messages.length > 0 ? ( + {readyState === ReadyState.CONNECTING && messages.length > 0 ? ( <> diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 5fba65a..d91c194 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - nav("/"); + window.location.href = "/"; }} > diff --git a/ta-browser/core/utils/openai_msg_parser.py b/ta-browser/core/utils/openai_msg_parser.py index 1458125..af57cdc 100644 --- a/ta-browser/core/utils/openai_msg_parser.py +++ b/ta-browser/core/utils/openai_msg_parser.py @@ -646,6 +646,7 @@ def get_conversation_history(self) -> List[ChatCompletionMessageParam]: return self.conversation_history +# In openai_msg_parser.py (or wherever ConversationStorage is defined) from datetime import datetime import os import json From cdfa7b59c82b02161f8e186757b0b65043bd57f6 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Tue, 22 Apr 2025 22:50:04 +0530 Subject: [PATCH 09/24] feat[openai compatibility]:made changes that CortexON can now switches between OpenAI and Antropic models --- cortex_on/agents/orchestrator_agent.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index c344213..ec2188e 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -160,7 +160,8 @@ class orchestrator_deps: """ async def orchestrator_agent(model_preference: str) -> Agent: - if model_preference == "Anthropic": + if model_preference == "Anthropic": + model = get_anthropic_model_instance() elif model_preference == "OpenAI": model = get_openai_model_instance() @@ -259,7 +260,8 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: # Create deps with the new stream_output deps_for_coder_agent = CoderAgentDeps( websocket=ctx.deps.websocket, - stream_output=coder_stream_output + stream_output=coder_stream_output, + model_preference=ctx.deps.model_preference ) # Run coder agent @@ -479,7 +481,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas planner_stream_output.steps.append("Asking planner to update the plan...") await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - updated_plan_response = await planner_agent.run(user_prompt=update_prompt) + agent = await planner_agent(model_preference=ctx.deps.model_preference) + updated_plan_response = await agent.run(user_prompt=update_prompt) updated_plan = updated_plan_response.data.plan @@ -506,7 +509,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.error(error_msg, exc_info=True) planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = a500 + planner_stream_output.status_code = 500 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) return f"Failed to update the plan: {error_msg}" From ea376dbaf11c09885f1114ad9741f34c72078b20 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Tue, 22 Apr 2025 23:31:41 +0530 Subject: [PATCH 10/24] feat[openai compatibility]:made changes that CortexON can now switches between OpenAI and Antropic models --- cortex_on/main.py | 12 +++++++++--- frontend/src/components/home/Header.tsx | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cortex_on/main.py b/cortex_on/main.py index bb3aa93..6d06755 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -7,6 +7,9 @@ from fastapi.middleware.cors import CORSMiddleware import logfire +# Configure Logfire +logfire.configure() + # Local application imports from instructor import SystemInstructor @@ -102,11 +105,14 @@ async def websocket_endpoint(websocket: WebSocket): logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) except WebSocketDisconnect: - print("Websocket disconnected") + print("[WEBSOCKET] Client disconnected") break except Exception as e: print(f"[WEBSOCKET] Error handling message: {str(e)}") + if "disconnect message has been received" in str(e): + print(f"[WEBSOCKET] DIsconnect detected, closing connection: {str(e)}") + break except Exception as e: print(f"[WEBSOCKET] Connection error: {str(e)}") - # finally: - # print("[WEBSOCKET] connection closed") + # finally: + # print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index d91c194..5fba65a 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - window.location.href = "/"; + nav("/"); }} > From bc066395a543937c77a284ac07af837b216ba5eb Mon Sep 17 00:00:00 2001 From: aryan Date: Wed, 23 Apr 2025 17:32:17 +0530 Subject: [PATCH 11/24] refactor(Header): update header layout and styling for improved responsiveness and usability - Adjusted logo container size and alignment - Enhanced button and model toggle button styling for better visibility - Improved overall layout structure for better responsiveness --- frontend/src/components/home/Header.tsx | 54 +++++++++++++------------ 1 file changed, 29 insertions(+), 25 deletions(-) diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 5fba65a..3c212fc 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -13,20 +13,20 @@ const Header = () => { return (
-
{ - dispatch(setMessages([])); - nav("/"); - }} - > - Logo -
-
+
{ + dispatch(setMessages([])); + nav("/"); + }} + > + Logo +
+
nav("/vault")} - className={`h-full flex justify-center items-center cursor-pointer border-b-2 px-4 hover:border-[#BD24CA] ${ + className={`h-full flex justify-center items-center cursor-pointer border-b-2 -mb-[2px] px-4 hover:border-[#BD24CA] ${ location.includes("/vault") ? "border-[#BD24CA]" : "border-background" @@ -34,22 +34,26 @@ const Header = () => { >

Vault

- - {/* Model Toggle Button */} -
- +
+ {/* Model Toggle Button - Smaller size */} +
+ +
+ + +
); }; From f929a683694340ffd23db6ed3a2410c331970508 Mon Sep 17 00:00:00 2001 From: aryan Date: Tue, 29 Apr 2025 21:33:51 +0530 Subject: [PATCH 12/24] refactor(logging): replace print statements with logfire for consistent logging across the application --- cortex_on/agents/orchestrator_agent.py | 3 --- cortex_on/agents/planner_agent.py | 1 - cortex_on/agents/web_surfer.py | 2 +- cortex_on/instructor.py | 14 -------------- cortex_on/main.py | 14 ++++---------- frontend/src/components/home/Header.tsx | 2 +- ta-browser/core/orchestrator.py | 8 +++----- ta-browser/core/utils/init_client.py | 7 ------- 8 files changed, 9 insertions(+), 42 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index ec2188e..e08c881 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -167,7 +167,6 @@ async def orchestrator_agent(model_preference: str) -> Agent: model = get_openai_model_instance() else: raise ValueError(f"Unknown model_preference: {model_preference}") - print(f"[ORCH_INIT] Creating orchestrator agent with model: {model}") orchestrator_agent = Agent( model=model, name="Orchestrator Agent", @@ -528,8 +527,6 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.info("All tools initialized for orchestrator agent") - print(f"[ORCH_INIT] Orchestrator agent initialized successfully: {orchestrator_agent}") - return orchestrator_agent diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index d5d7832..d607d04 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -183,7 +183,6 @@ async def planner_agent(model_preference: str = "Anthropic") -> Agent: model = get_openai_model_instance() else: raise ValueError(f"Unknown model_preference: {model_preference}") - print(f"[PLANNER_INIT] Creating planner agent with model: {model}") planner_agent = Agent( model=model, name="Planner Agent", diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index e9fe772..2711348 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -47,7 +47,7 @@ async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, An "critique_disabled": False, "model_preference": self.model_preference } - logfire.info(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") + logfire.info(f"Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index 388a7bb..bf9670f 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -25,11 +25,6 @@ load_dotenv() -print("[INIT] Loading agents...") -print(f"[INIT] Planner agent loaded: {planner_agent}") -print(f"[INIT] Coder agent loaded: {coder_agent}") -print(f"[INIT] Orchestrator agent loaded: {orchestrator_agent}") - class DateTimeEncoder(json.JSONEncoder): """Custom JSON encoder that can handle datetime objects""" def default(self, obj): @@ -47,7 +42,6 @@ def __init__(self, model_preference: str = "Anthropic"): self.orchestrator_response: List[StreamResponse] = [] self.model_preference = model_preference self._setup_logging() - print("[INIT] SystemInstructor initialization complete") def _setup_logging(self) -> None: """Configure logging with proper formatting""" @@ -94,8 +88,6 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: await self._safe_websocket_send(stream_output) stream_output.steps.append("Agents initialized successfully") await self._safe_websocket_send(stream_output) - - print("Orchestrator agent running") agent = await orchestrator_agent(self.model_preference) orchestrator_response = await agent.run( @@ -133,12 +125,6 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: async def shutdown(self): """Clean shutdown of orchestrator""" try: - #Close websocket if open - # if self.websocket: - # await self.websocket.close() - # logfire.info("websocket shutdown complete") - # print("websocket shutdown complete") - # Clear all responses self.orchestrator_response = [] logfire.info("Orchestrator shutdown complete") diff --git a/cortex_on/main.py b/cortex_on/main.py index 6d06755..44b840b 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -28,13 +28,9 @@ async def lifespan(app: FastAPI): # Initialize the instructor global instructor instructor = SystemInstructor(model_preference=MODEL_PREFERENCE) - print("[STARTUP] Instructor initialized") yield - - # # Cleanup - # if instructor: - # await instructor.shutdown() + app: FastAPI = FastAPI(lifespan=lifespan) @@ -105,14 +101,12 @@ async def websocket_endpoint(websocket: WebSocket): logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) except WebSocketDisconnect: - print("[WEBSOCKET] Client disconnected") + logfire.info("[WEBSOCKET] Client disconnected") break except Exception as e: - print(f"[WEBSOCKET] Error handling message: {str(e)}") + logfire.error(f"[WEBSOCKET] Error handling message: {str(e)}") if "disconnect message has been received" in str(e): print(f"[WEBSOCKET] DIsconnect detected, closing connection: {str(e)}") break except Exception as e: - print(f"[WEBSOCKET] Connection error: {str(e)}") - # finally: - # print("[WEBSOCKET] connection closed") + logfire.error(f"[WEBSOCKET] Connection error: {str(e)}") diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 3c212fc..4d142fb 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -26,7 +26,7 @@ const Header = () => {
nav("/vault")} - className={`h-full flex justify-center items-center cursor-pointer border-b-2 -mb-[2px] px-4 hover:border-[#BD24CA] ${ + className={`h-full flex justify-center items-center cursor-pointer border-b-2 -mb-[27px] px-4 pb-7 hover:border-[#BD24CA] ${ location.includes("/vault") ? "border-[#BD24CA]" : "border-background" diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 69df1d7..4d56511 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -252,7 +252,7 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com", m "start_url": start_url, "model_preference": model_preference }) - print(f"[ORCHESTRATOR] Starting async_init with model_preference: {model_preference}") + logger.info(f"Starting async_init with model_preference: {model_preference}") # Store and validate parameters self.job_id = job_id @@ -283,15 +283,13 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com", m # 4. Initialize client and agents try: - print(f"[ORCHESTRATOR] Calling initialize_client with model_preference: {model_preference}") from core.utils.init_client import initialize_client self.client, model_instance = await initialize_client(model_preference) - print(f"[ORCHESTRATOR] Client initialized successfully, now initializing agents") self.initialize_agents(model_instance) - print(f"[ORCHESTRATOR] Agents initialized successfully with model_preference: {model_preference}") + logger.info(f"Agents initialized successfully with model_preference: {model_preference}") except Exception as agent_error: error_msg = f"Failed to initialize client and agents: {str(agent_error)}" - print(f"[ORCHESTRATOR] ERROR: {error_msg}") + logger.error(f"{error_msg}") raise RuntimeError(error_msg) from agent_error self.async_init_done = True diff --git a/ta-browser/core/utils/init_client.py b/ta-browser/core/utils/init_client.py index 5a423b9..5be2ea3 100644 --- a/ta-browser/core/utils/init_client.py +++ b/ta-browser/core/utils/init_client.py @@ -16,7 +16,6 @@ async def initialize_client(model_preference: str = "Anthropic"): tuple: (client_instance, model_instance) """ try: - print(f"[INIT_CLIENT] *** Initializing client with model preference: {model_preference} ***") logger.info(f"Initializing client with model preference: {model_preference}") if model_preference == "Anthropic": # Get API key from environment variable @@ -27,7 +26,6 @@ async def initialize_client(model_preference: str = "Anthropic"): # Set model name - Claude 3.5 Sonnet model_name = os.getenv("ANTHROPIC_MODEL_NAME") - print(f"[INIT_CLIENT] Using Anthropic model: {model_name}") # Create client config config = { @@ -44,7 +42,6 @@ async def initialize_client(model_preference: str = "Anthropic"): from pydantic_ai.models.anthropic import AnthropicModel model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) - print(f"[INIT_CLIENT] Anthropic client initialized successfully with model: {model_name}") logger.info(f"Anthropic client initialized successfully with model: {model_name}") return client_instance, model_instance elif model_preference == "OpenAI": @@ -56,7 +53,6 @@ async def initialize_client(model_preference: str = "Anthropic"): # Set model name - GPT-4o model_name = os.getenv("OPENAI_MODEL_NAME") - print(f"[INIT_CLIENT] Using OpenAI model: {model_name}") # Create client config config = { @@ -73,16 +69,13 @@ async def initialize_client(model_preference: str = "Anthropic"): from pydantic_ai.models.openai import OpenAIModel model_instance = OpenAIModel(model_name=model_name, openai_client=client_instance) - print(f"[INIT_CLIENT] OpenAI client initialized successfully with model: {model_name}") logger.info(f"OpenAI client initialized successfully with model: {model_name}") return client_instance, model_instance else: error_msg = f"Invalid model preference: {model_preference}. Must be 'Anthropic' or 'OpenAI'" - print(f"[INIT_CLIENT] ERROR: {error_msg}") raise ValueError(error_msg) except Exception as e: error_msg = f"Error initializing client: {str(e)}" - print(f"[INIT_CLIENT] CRITICAL ERROR: {error_msg}") logger.error(error_msg, exc_info=True) raise \ No newline at end of file From 194735faefae1ef8f6ef8c6f8606dddb5c7e9697 Mon Sep 17 00:00:00 2001 From: aryan Date: Mon, 14 Apr 2025 00:58:05 +0530 Subject: [PATCH 13/24] feat(model-preference): Implement model preference selection across the application - Added model preference handling in SystemInstructor and WebSurfer classes. - Introduced API endpoints for setting and retrieving model preference. - Updated FastAPI app to manage model preference during websocket connections and response generation. - Enhanced orchestrator and session management to utilize model preference. - Integrated model preference into frontend components for user selection. --- cortex_on/agents/orchestrator_agent.py | 8 +- cortex_on/agents/web_surfer.py | 10 +- cortex_on/instructor.py | 6 +- cortex_on/main.py | 62 +++++++++-- frontend/src/components/home/Header.tsx | 12 ++- frontend/src/components/ui/ModelToggle.tsx | 61 +++++++++++ .../src/dataStore/modelPreferenceSlice.ts | 22 ++++ frontend/src/dataStore/store.ts | 7 +- frontend/src/lib/utils.ts | 4 +- frontend/src/services/modelPreferenceApi.ts | 23 ++++ .../task_20250413_190916/conversation.json | 79 ++++++++++++++ .../task_20250413_191154/conversation.json | 1 + .../task_20250413_192014/conversation.json | 49 +++++++++ ta-browser/core/orchestrator.py | 52 ++++----- ta-browser/core/server/models/web.py | 1 + ta-browser/core/server/routes/web.py | 7 +- .../core/server/utils/session_tracker.py | 11 +- ta-browser/core/utils/init_client.py | 102 +++++++++++++----- ta-browser/core/utils/openai_client.py | 65 +---------- 19 files changed, 441 insertions(+), 141 deletions(-) create mode 100644 frontend/src/components/ui/ModelToggle.tsx create mode 100644 frontend/src/dataStore/modelPreferenceSlice.ts create mode 100644 frontend/src/services/modelPreferenceApi.ts create mode 100644 ta-browser/1744571356739037821/task_20250413_190916/conversation.json create mode 100644 ta-browser/1744571514482277604/task_20250413_191154/conversation.json create mode 100644 ta-browser/1744572014568045555/task_20250413_192014/conversation.json diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 6b001ba..284d927 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -22,6 +22,7 @@ class orchestrator_deps: stream_output: Optional[StreamResponse] = None # Add a collection to track agent-specific streams agent_responses: Optional[List[StreamResponse]] = None + model_preference: str = "Anthropic" orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. @@ -307,8 +308,11 @@ async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - # Initialize WebSurfer agent - web_surfer_agent = WebSurfer(api_url="http://localhost:8000/api/v1/web/stream") + # Initialize WebSurfer agent with model preference + web_surfer_agent = WebSurfer( + api_url="http://localhost:8000/api/v1/web/stream", + model_preference=ctx.deps.model_preference + ) # Run WebSurfer with its own stream_output success, message, messages = await web_surfer_agent.generate_reply( diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 34e2cfd..03503d5 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -29,19 +29,25 @@ TIMEOUT = 9999999999999999999999999999999999999999999 class WebSurfer: - def __init__(self, api_url: str = "http://localhost:8000/api/v1/web/stream"): + def __init__(self, api_url: str = "http://localhost:8000/api/v1/web/stream", model_preference: str = "Anthropic"): self.api_url = api_url self.name = "Web Surfer Agent" self.description = "An agent that is a websurfer and a webscraper that can access any web-page to extract information or perform actions." self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None + self.model_preference = model_preference async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, Any]]]: session_timeout = aiohttp.ClientTimeout(total=None, sock_connect=TIMEOUT, sock_read=TIMEOUT) async with aiohttp.ClientSession(timeout=session_timeout) as session: final_json_response = [] try: - payload = {"cmd": instruction, "critique_disabled": False} + payload = { + "cmd": instruction, + "critique_disabled": False, + "model_preference": self.model_preference + } + print(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index b4f0efb..f2688b4 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -38,10 +38,11 @@ def default(self, obj): # Main Orchestrator Class class SystemInstructor: - def __init__(self): + def __init__(self, model_preference: str = "Anthropic"): self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None self.orchestrator_response: List[StreamResponse] = [] + self.model_preference = model_preference self._setup_logging() def _setup_logging(self) -> None: @@ -80,7 +81,8 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: deps_for_orchestrator = orchestrator_deps( websocket=self.websocket, stream_output=stream_output, - agent_responses=self.orchestrator_response # Pass reference to collection + agent_responses=self.orchestrator_response, # Pass reference to collection + model_preference=self.model_preference ) try: diff --git a/cortex_on/main.py b/cortex_on/main.py index a8dd4de..e1adddd 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -1,27 +1,75 @@ # Standard library imports from typing import List, Optional +from contextlib import asynccontextmanager # Third-party imports -from fastapi import FastAPI, WebSocket +from fastapi import FastAPI, WebSocket, Depends +from fastapi.middleware.cors import CORSMiddleware # Local application imports from instructor import SystemInstructor -app: FastAPI = FastAPI() +# Default model preference is Anthropic +MODEL_PREFERENCE = "Anthropic" -async def generate_response(task: str, websocket: Optional[WebSocket] = None): - orchestrator: SystemInstructor = SystemInstructor() +@asynccontextmanager +async def lifespan(app: FastAPI): + # Set default model preference at startup + app.state.model_preference = MODEL_PREFERENCE + print(f"[STARTUP] Setting default model preference to: {MODEL_PREFERENCE}") + yield + + +app: FastAPI = FastAPI(lifespan=lifespan) + +# Add CORS middleware +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], # Allow all origins for development + allow_credentials=True, + allow_methods=["*"], # Allow all methods + allow_headers=["*"], # Allow all headers +) + +async def get_model_preference() -> str: + """ + Get the current model preference from app state + """ + print(f"[GET] Current model preference: {app.state.model_preference}") + return app.state.model_preference + +@app.get("/set_model_preference") +async def set_model_preference(model: str): + """ + Set the model preference (Anthropic or OpenAI) + """ + if model not in ["Anthropic", "OpenAI"]: + print(f"[ERROR] Invalid model preference attempted: {model}") + return {"error": "Model must be 'Anthropic' or 'OpenAI'"} + print(f"[SET] Changing model preference from {app.state.model_preference} to {model}") + app.state.model_preference = model + return {"message": f"Model preference set to {model}"} + +async def generate_response(task: str, websocket: Optional[WebSocket] = None, model_preference: str = None): + if model_preference is None: + model_preference = app.state.model_preference + print(f"[GENERATE] Using model preference: {model_preference} for task: {task[:30]}...") + orchestrator: SystemInstructor = SystemInstructor(model_preference) return await orchestrator.run(task, websocket) @app.get("/agent/chat") -async def agent_chat(task: str) -> List: - final_agent_response = await generate_response(task) +async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List: + print(f"[CHAT] Received chat request with model preference: {model_preference}") + final_agent_response = await generate_response(task, model_preference=model_preference) return final_agent_response @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() + model_preference = app.state.model_preference + print(f"[WEBSOCKET] New connection using model preference: {model_preference}") while True: data = await websocket.receive_text() - await generate_response(data, websocket) + print(f"[WEBSOCKET] Received message, using model: {model_preference}") + await generate_response(data, websocket, model_preference) diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index d6217fd..5fba65a 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -4,6 +4,7 @@ import {useDispatch} from "react-redux"; import {useLocation, useNavigate} from "react-router-dom"; import Logo from "../../assets/CortexON_logo_dark.svg"; import {Button} from "../ui/button"; +import ModelToggle from "../ui/ModelToggle"; const Header = () => { const nav = useNavigate(); @@ -13,7 +14,7 @@ const Header = () => { return (
{ dispatch(setMessages([])); nav("/"); @@ -21,10 +22,11 @@ const Header = () => { > Logo
-
+ +
nav("/vault")} - className={`w-[10%] h-full flex justify-center items-center cursor-pointer border-b-2 hover:border-[#BD24CA] ${ + className={`h-full flex justify-center items-center cursor-pointer border-b-2 px-4 hover:border-[#BD24CA] ${ location.includes("/vault") ? "border-[#BD24CA]" : "border-background" @@ -32,7 +34,11 @@ const Header = () => { >

Vault

+ + {/* Model Toggle Button */} +
+ + + + GPT-4o + +
+ ); +} \ No newline at end of file diff --git a/frontend/src/dataStore/modelPreferenceSlice.ts b/frontend/src/dataStore/modelPreferenceSlice.ts new file mode 100644 index 0000000..7c1d602 --- /dev/null +++ b/frontend/src/dataStore/modelPreferenceSlice.ts @@ -0,0 +1,22 @@ +import { createSlice, PayloadAction } from "@reduxjs/toolkit"; + +interface ModelPreferenceState { + preference: "Anthropic" | "OpenAI"; +} + +const initialState: ModelPreferenceState = { + preference: "Anthropic", +}; + +const modelPreferenceSlice = createSlice({ + name: "modelPreference", + initialState, + reducers: { + setModelPreference: (state, action: PayloadAction<"Anthropic" | "OpenAI">) => { + state.preference = action.payload; + }, + }, +}); + +export const { setModelPreference } = modelPreferenceSlice.actions; +export default modelPreferenceSlice; \ No newline at end of file diff --git a/frontend/src/dataStore/store.ts b/frontend/src/dataStore/store.ts index c3be4e6..516535f 100644 --- a/frontend/src/dataStore/store.ts +++ b/frontend/src/dataStore/store.ts @@ -1,13 +1,18 @@ import vaultApi from "@/services/vaultApi"; +import modelPreferenceApi from "@/services/modelPreferenceApi"; import {configureStore} from "@reduxjs/toolkit"; import messagesSlice from "./messagesSlice"; +import modelPreferenceSlice from "./modelPreferenceSlice"; + export const store = configureStore({ reducer: { [messagesSlice.name]: messagesSlice.reducer, [vaultApi.reducerPath]: vaultApi.reducer, + [modelPreferenceSlice.name]: modelPreferenceSlice.reducer, + [modelPreferenceApi.reducerPath]: modelPreferenceApi.reducer, }, middleware: (getDefaultMiddleware) => - getDefaultMiddleware().concat(vaultApi.middleware), + getDefaultMiddleware().concat(vaultApi.middleware, modelPreferenceApi.middleware), }); export type RootState = ReturnType; diff --git a/frontend/src/lib/utils.ts b/frontend/src/lib/utils.ts index dfe0774..07ab4f7 100644 --- a/frontend/src/lib/utils.ts +++ b/frontend/src/lib/utils.ts @@ -1,5 +1,5 @@ -import {clsx, type ClassValue} from "clsx"; -import {twMerge} from "tailwind-merge"; +import { type ClassValue, clsx } from "clsx"; +import { twMerge } from "tailwind-merge"; export function cn(...inputs: ClassValue[]) { return twMerge(clsx(inputs)); diff --git a/frontend/src/services/modelPreferenceApi.ts b/frontend/src/services/modelPreferenceApi.ts new file mode 100644 index 0000000..e7320a6 --- /dev/null +++ b/frontend/src/services/modelPreferenceApi.ts @@ -0,0 +1,23 @@ +import { createApi } from "@reduxjs/toolkit/query/react"; +import { fetchBaseQuery } from "@reduxjs/toolkit/query/react"; + +// Get the base URL from environment variable or use a default +const CORTEX_ON_API_URL = import.meta.env.VITE_CORTEX_ON_API_URL || "http://localhost:8081"; + +const modelPreferenceApi = createApi({ + reducerPath: "modelPreferenceApi", + baseQuery: fetchBaseQuery({ + baseUrl: CORTEX_ON_API_URL, + }), + endpoints: (builder) => ({ + setModelPreference: builder.mutation<{ message: string }, string>({ + query: (model) => ({ + url: `/set_model_preference?model=${model}`, + method: "GET", + }), + }), + }), +}); + +export const { useSetModelPreferenceMutation } = modelPreferenceApi; +export default modelPreferenceApi; \ No newline at end of file diff --git a/ta-browser/1744571356739037821/task_20250413_190916/conversation.json b/ta-browser/1744571356739037821/task_20250413_190916/conversation.json new file mode 100644 index 0000000..e9976ca --- /dev/null +++ b/ta-browser/1744571356739037821/task_20250413_190916/conversation.json @@ -0,0 +1,79 @@ +[ + { + "role": "user", + "content": "Inputs :\nUser Query : Go to Amazon.com and search for RTX 4060 graphics cards. Return the search results including prices and available models.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" + }, + { + "role": "assistant", + "content": "Plan: 1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\nNext Step: Navigate to Amazon's website (https://www.amazon.com)", + "name": "planner_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" current_step=\"Navigate to Amazon's website (https://www.amazon.com)\" " + }, + { + "role": "assistant", + "content": "I'll help navigate to Amazon's website using the open_url_tool.", + "name": "browser_nav_agent" + }, + { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", + "type": "function", + "function": { + "name": "open_url_tool", + "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", + "name": "open_url_tool", + "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" + }, + { + "role": "assistant", + "content": "I attempted to navigate to Amazon's website but encountered an error. This is common with Amazon as they have strict bot detection measures. Let me verify if we landed on any page by checking the current URL.", + "name": "browser_nav_agent" + }, + { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", + "type": "function", + "function": { + "name": "get_url", + "arguments": "{\"raw_args\": \"ArgsDict(args_dict={})\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", + "name": "get_url", + "content": "Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more." + }, + { + "role": "assistant", + "content": "Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.", + "name": "browser_nav_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" next_step=\"Navigate to Amazon's website (https://www.amazon.com)\" tool_response=\"Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\nTool Call: get_url\nArguments: {}\nResponse: Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more.\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" + }, + { + "role": "assistant", + "content": "Feedback: Original Plan:\n1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\n\nCurrent Progress:\nCurrently on Step 1: Navigate to Amazon's website (https://www.amazon.com)\n\nAnalysis and Feedback:\nThe browser agent has successfully completed the first step of navigating to Amazon.com. This is confirmed by:\n1. The get_url tool response showing we are at \"https://www.amazon.com/\"\n2. The page title verification showing \"Amazon.com. Spend less. Smile more.\"\n\nDespite an initial execution context error (which is common during page loads), the navigation was ultimately successful. The browser agent can now proceed to the next step of locating the search bar on Amazon's homepage.\n\nThe plan should continue with Step 2 as we have successfully completed Step 1.\nTerminate: False\nFinal Response: None", + "name": "critique_agent" + } +] \ No newline at end of file diff --git a/ta-browser/1744571514482277604/task_20250413_191154/conversation.json b/ta-browser/1744571514482277604/task_20250413_191154/conversation.json new file mode 100644 index 0000000..0637a08 --- /dev/null +++ b/ta-browser/1744571514482277604/task_20250413_191154/conversation.json @@ -0,0 +1 @@ +[] \ No newline at end of file diff --git a/ta-browser/1744572014568045555/task_20250413_192014/conversation.json b/ta-browser/1744572014568045555/task_20250413_192014/conversation.json new file mode 100644 index 0000000..d214df6 --- /dev/null +++ b/ta-browser/1744572014568045555/task_20250413_192014/conversation.json @@ -0,0 +1,49 @@ +[ + { + "role": "user", + "content": "Inputs :\nUser Query : Go to Amazon website and search for \"rtx 4060\". Extract the product listings including prices, models, and ratings for RTX 4060 graphics cards.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" + }, + { + "role": "assistant", + "content": "Plan: 1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\nNext Step: Open Amazon's website via direct URL: https://www.amazon.com", + "name": "planner_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" current_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" " + }, + { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", + "type": "function", + "function": { + "name": "open_url_tool", + "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", + "name": "open_url_tool", + "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" + }, + { + "role": "assistant", + "content": "I encountered an error while trying to open Amazon's website. The page failed to load successfully.", + "name": "browser_nav_agent" + }, + { + "role": "user", + "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" next_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" tool_response=\"I encountered an error while trying to open Amazon's website. The page failed to load successfully.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" + }, + { + "role": "assistant", + "content": "Feedback: The original plan involves opening Amazon's website and then performing a series of actions to search for 'RTX 4060' and extract product details. The current step was to open Amazon's website using the direct URL. However, the browser agent encountered an error while trying to load the page. The error message indicates that the execution context was destroyed, likely due to a navigation issue. This means that the page could not be successfully loaded, and thus we cannot proceed with the search and extraction tasks. Since this is a critical failure at the very first step, it's not possible to continue with the plan.\nTerminate: True\nFinal Response: The task encountered an error while trying to open Amazon's website, which prevented further actions. The error was related to the execution context being destroyed, likely due to a navigation issue. This means the page could not be loaded successfully, and as a result, we cannot proceed with searching for RTX 4060 graphics cards or extracting any product information. I recommend retrying the task or checking the network connection and URL for any issues.", + "name": "critique_agent" + } +] \ No newline at end of file diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 9dbf16e..6a2aaef 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -244,27 +244,32 @@ def log_token_usage(self, agent_type: str, usage: Usage, step: Optional[int] = N """ ) - async def async_init(self, job_id: str, start_url: str = "https://google.com"): + async def async_init(self, job_id: str, start_url: str = "https://google.com", model_preference: str = "Anthropic"): """Initialize a new session context with improved error handling""" try: logger.info("Initializing browser session", extra={ "job_id": job_id, - "start_url": start_url + "start_url": start_url, + "model_preference": model_preference }) - - # 1. Initialize job_id and validate + print(f"[ORCHESTRATOR] Starting async_init with model_preference: {model_preference}") + + # Store and validate parameters + self.job_id = job_id + if not job_id: - raise ValueError("job_id is required for initialization") - self.job_id = str(job_id) - logger.debug(f"job_id: {self.job_id}") + raise ValueError("job_id is required") - # 2. Initialize conversation storage + # 1. Create conversation_storage with job_id + self.conversation_storage = ConversationStorage(job_id) + + # 2. Initialize browser manager try: - self.conversation_storage = ConversationStorage(job_id=self.job_id) - except Exception as storage_error: - raise RuntimeError(f"Failed to initialize conversation storage: {str(storage_error)}") from storage_error - - # 3. Set and validate URL + await self.initialize_browser_manager(start_url) + except Exception as browser_error: + raise RuntimeError(f"Failed to initialize browser: {str(browser_error)}") from browser_error + + # 3. Set current URL and domain try: self.current_url = start_url self.current_domain = extract_domain(self.current_url) @@ -273,23 +278,18 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com"): except InvalidURLError as url_error: raise ValueError(f"Invalid URL provided: {str(url_error)}") from url_error - # 4. Initialize browser manager with start_url - if not self.browser_manager: - try: - self.browser_manager = await self.initialize_browser_manager(start_url=start_url) - if not self.browser_manager: - raise RuntimeError("Browser manager initialization failed") - except Exception as browser_error: - raise RuntimeError(f"Failed to initialize browser manager: {str(browser_error)}") from browser_error - - # 5. Initialize client and agents + # 4. Initialize client and agents try: - + print(f"[ORCHESTRATOR] Calling initialize_client with model_preference: {model_preference}") from core.utils.init_client import initialize_client - self.client, model_instance = await initialize_client() + self.client, model_instance = await initialize_client(model_preference) + print(f"[ORCHESTRATOR] Client initialized successfully, now initializing agents") self.initialize_agents(model_instance) + print(f"[ORCHESTRATOR] Agents initialized successfully with model_preference: {model_preference}") except Exception as agent_error: - raise RuntimeError(f"Failed to initialize client and agents: {str(agent_error)}") from agent_error + error_msg = f"Failed to initialize client and agents: {str(agent_error)}" + print(f"[ORCHESTRATOR] ERROR: {error_msg}") + raise RuntimeError(error_msg) from agent_error self.async_init_done = True logger.debug("Async initialization completed successfully") diff --git a/ta-browser/core/server/models/web.py b/ta-browser/core/server/models/web.py index bf9e9fc..ca397f0 100644 --- a/ta-browser/core/server/models/web.py +++ b/ta-browser/core/server/models/web.py @@ -7,6 +7,7 @@ class StreamRequestModel(BaseModel): cmd: str = Field(..., description="Command to execute") url: str = Field("https://google.com", description="URL to navigate to") critique_disabled: bool = Field(False, description="Whether to disable critique") + model_preference: str = Field("Anthropic", description="Model provider preference (Anthropic or OpenAI)") @validator('url') def validate_and_format_url(cls, v): diff --git a/ta-browser/core/server/routes/web.py b/ta-browser/core/server/routes/web.py index a5cee6d..0f0ee09 100644 --- a/ta-browser/core/server/routes/web.py +++ b/ta-browser/core/server/routes/web.py @@ -115,6 +115,7 @@ async def stream_session( and streams back real-time updates as the command is executed. """ print(f"[{time.time()}] Stream route: Starting processing") + print(f"[STREAM] Received request with model_preference: {request.model_preference}") session_tracker = SessionTracker() @@ -124,8 +125,12 @@ async def stream_session( # Initialize the session logger.debug(f"Initializing stream session with ID {session_id}") + print(f"[STREAM] Initializing session with model_preference: {request.model_preference}") session_info = await session_tracker.initialize_session( - request.url, request.critique_disabled, session_id + request.url, + request.critique_disabled, + session_id, + request.model_preference ) session_context = session_tracker.active_sessions.get(session_id) diff --git a/ta-browser/core/server/utils/session_tracker.py b/ta-browser/core/server/utils/session_tracker.py index acf7168..fb3fa13 100644 --- a/ta-browser/core/server/utils/session_tracker.py +++ b/ta-browser/core/server/utils/session_tracker.py @@ -28,17 +28,19 @@ def get_active_sessions_status(self) -> dict: "sessions": list(self.active_sessions.keys()) } - async def initialize_session(self, start_url: str, no_crit: bool, session_id: str) -> Dict[str, Any]: + async def initialize_session(self, start_url: str, no_crit: bool, session_id: str, model_preference: str = "Anthropic") -> Dict[str, Any]: """Initialize a new session context""" logger.debug(f"Starting session initialization with URL: {start_url}") + print(f"[SESSION_TRACKER] Initializing session with model_preference: {model_preference}") logger.set_job_id(session_id) orchestrator = None try: orchestrator = Orchestrator(input_mode="API", no_crit=no_crit) - await orchestrator.async_init(job_id=session_id, start_url=start_url) - logger.debug(f"Orchestrator async_init completed with params: job_id={session_id}, start_url={start_url}") + print(f"[SESSION_TRACKER] Passing model_preference: {model_preference} to orchestrator.async_init") + await orchestrator.async_init(job_id=session_id, start_url=start_url, model_preference=model_preference) + logger.debug(f"Orchestrator async_init completed with params: job_id={session_id}, start_url={start_url}, model_preference={model_preference}") notification_queue = Queue() orchestrator.notification_queue = notification_queue @@ -49,7 +51,8 @@ async def initialize_session(self, start_url: str, no_crit: bool, session_id: st "notification_queue": notification_queue, "start_time": datetime.now(), "current_url": start_url, - "include_screenshot": False + "include_screenshot": False, + "model_preference": model_preference } self.add_active_session(session_id, session_context) diff --git a/ta-browser/core/utils/init_client.py b/ta-browser/core/utils/init_client.py index 7d170c6..5a423b9 100644 --- a/ta-browser/core/utils/init_client.py +++ b/ta-browser/core/utils/init_client.py @@ -1,44 +1,88 @@ import os from core.utils.anthropic_client import create_client_with_retry as create_anthropic_client, AsyncAnthropic from core.utils.logger import Logger +from core.utils.openai_client import get_client logger = Logger() -async def initialize_client(): +async def initialize_client(model_preference: str = "Anthropic"): """ Initialize and return the Anthropic client and model instance + Args: + model_preference (str): The model provider to use ("Anthropic" or "OpenAI") + Returns: tuple: (client_instance, model_instance) """ try: - # Get API key from environment variable - api_key = os.getenv("ANTHROPIC_API_KEY") - if not api_key: - logger.error("ANTHROPIC_API_KEY not found in environment variables") - raise ValueError("ANTHROPIC_API_KEY not found in environment variables") - - # Set model name - Claude 3.5 Sonnet - model_name = os.getenv("ANTHROPIC_MODEL_NAME") - - # Create client config - config = { - "api_key": api_key, - "model": model_name, - "max_retries": 3, - "timeout": 300.0 - } - - # Initialize client - client_instance = create_anthropic_client(AsyncAnthropic, config) - - # Create model instance - from pydantic_ai.models.anthropic import AnthropicModel - model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) - - logger.info(f"Anthropic client initialized successfully with model: {model_name}") - return client_instance, model_instance - + print(f"[INIT_CLIENT] *** Initializing client with model preference: {model_preference} ***") + logger.info(f"Initializing client with model preference: {model_preference}") + if model_preference == "Anthropic": + # Get API key from environment variable + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + logger.error("ANTHROPIC_API_KEY not found in environment variables") + raise ValueError("ANTHROPIC_API_KEY not found in environment variables") + + # Set model name - Claude 3.5 Sonnet + model_name = os.getenv("ANTHROPIC_MODEL_NAME") + print(f"[INIT_CLIENT] Using Anthropic model: {model_name}") + + # Create client config + config = { + "api_key": api_key, + "model": model_name, + "max_retries": 3, + "timeout": 300.0 + } + + # Initialize client + client_instance = create_anthropic_client(AsyncAnthropic, config) + + # Create model instance + from pydantic_ai.models.anthropic import AnthropicModel + model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) + + print(f"[INIT_CLIENT] Anthropic client initialized successfully with model: {model_name}") + logger.info(f"Anthropic client initialized successfully with model: {model_name}") + return client_instance, model_instance + elif model_preference == "OpenAI": + # Get API key from environment variable + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + logger.error("OPENAI_API_KEY not found in environment variables") + raise ValueError("OPENAI_API_KEY not found in environment variables") + + # Set model name - GPT-4o + model_name = os.getenv("OPENAI_MODEL_NAME") + print(f"[INIT_CLIENT] Using OpenAI model: {model_name}") + + # Create client config + config = { + "api_key": api_key, + "model": model_name, + "max_retries": 3, + "timeout": 300.0 + } + + # Initialize client + client_instance = get_client() + + # Create model instance + from pydantic_ai.models.openai import OpenAIModel + model_instance = OpenAIModel(model_name=model_name, openai_client=client_instance) + + print(f"[INIT_CLIENT] OpenAI client initialized successfully with model: {model_name}") + logger.info(f"OpenAI client initialized successfully with model: {model_name}") + return client_instance, model_instance + else: + error_msg = f"Invalid model preference: {model_preference}. Must be 'Anthropic' or 'OpenAI'" + print(f"[INIT_CLIENT] ERROR: {error_msg}") + raise ValueError(error_msg) + except Exception as e: - logger.error(f"Error initializing Anthropic client: {str(e)}", exc_info=True) + error_msg = f"Error initializing client: {str(e)}" + print(f"[INIT_CLIENT] CRITICAL ERROR: {error_msg}") + logger.error(error_msg, exc_info=True) raise \ No newline at end of file diff --git a/ta-browser/core/utils/openai_client.py b/ta-browser/core/utils/openai_client.py index 80670e3..5c74992 100644 --- a/ta-browser/core/utils/openai_client.py +++ b/ta-browser/core/utils/openai_client.py @@ -28,57 +28,15 @@ def validate_model(model: str) -> bool: return True @staticmethod def get_text_config() -> Dict: - model = get_env_var("AGENTIC_BROWSER_TEXT_MODEL") - if not OpenAIConfig.validate_model(model): - raise ModelValidationError( - f"Invalid model: {model}. Must match one of the patterns: " - f"{', '.join(OpenAIConfig.VALID_MODEL_PATTERNS)}" - ) - return { - "api_key": get_env_var("AGENTIC_BROWSER_TEXT_API_KEY"), - "base_url": get_env_var("AGENTIC_BROWSER_TEXT_BASE_URL"), - "model": model, + "api_key": get_env_var("OPENAI_API_KEY"), + "base_url": "https://api.openai.com/v1", + "model": get_env_var("OPENAI_MODEL_NAME"), "max_retries": 3, "timeout": 300.0 } - @staticmethod - def get_ss_config() -> Dict: - model = get_env_var("AGENTIC_BROWSER_SS_MODEL") - if not OpenAIConfig.validate_model(model): - raise ModelValidationError( - f"Invalid model: {model}. Must match one of the patterns: " - f"{', '.join(OpenAIConfig.VALID_MODEL_PATTERNS)}" - ) - - return { - "api_key": get_env_var("AGENTIC_BROWSER_SS_API_KEY"), - "base_url": get_env_var("AGENTIC_BROWSER_SS_BASE_URL"), - "model": model, - "max_retries": 3, - "timeout": 300.0 - } -async def validate_models(client: AsyncOpenAI) -> bool: - """Validate that configured models are available""" - try: - available_models = await client.models.list() - available_model_ids = [model.id for model in available_models.data] - - text_model = get_text_model() - ss_model = get_ss_model() - - if text_model not in available_model_ids: - raise ModelValidationError(f"Text model '{text_model}' not available. Available models: {', '.join(available_model_ids)}") - - if ss_model not in available_model_ids: - raise ModelValidationError(f"Screenshot model '{ss_model}' not available. Available models: {', '.join(available_model_ids)}") - - return True - except Exception as e: - logger.error(f"Model validation failed: {str(e)}") - return False def create_client_with_retry(client_class, config: dict): """Create an OpenAI client with proper error handling""" @@ -102,26 +60,9 @@ def get_client(): config = OpenAIConfig.get_text_config() return create_client_with_retry(AsyncOpenAI, config) -def get_ss_client(): - """Get OpenAI client for screenshot analysis""" - config = OpenAIConfig.get_ss_config() - return create_client_with_retry(OpenAI, config) - -def get_text_model() -> str: - """Get model name for text analysis""" - return OpenAIConfig.get_text_config()["model"] - -def get_ss_model() -> str: - """Get model name for screenshot analysis""" - return OpenAIConfig.get_ss_config()["model"] - # Example usage async def initialize_and_validate(): """Initialize client and validate configuration""" client = get_client() - # Validate models - if not await validate_models(client): - raise ModelValidationError("Failed to validate models. Please check your configuration.") - return client \ No newline at end of file From 5e9130816c84e836288be2a67421f2bf94ccf8c6 Mon Sep 17 00:00:00 2001 From: aryan Date: Mon, 14 Apr 2025 16:47:23 +0530 Subject: [PATCH 14/24] fix(conversation_handling): Remove obsolete conversation files and enhance error handling in orchestrator. - Deleted unused conversation JSON files - Improved error handling during conversation storage initialization in the orchestrator to provide clearer feedback on failures. --- .../task_20250413_190916/conversation.json | 79 ------------------- .../task_20250413_191154/conversation.json | 1 - .../task_20250413_192014/conversation.json | 49 ------------ ta-browser/core/orchestrator.py | 5 +- ta-browser/core/utils/openai_msg_parser.py | 1 - 5 files changed, 4 insertions(+), 131 deletions(-) delete mode 100644 ta-browser/1744571356739037821/task_20250413_190916/conversation.json delete mode 100644 ta-browser/1744571514482277604/task_20250413_191154/conversation.json delete mode 100644 ta-browser/1744572014568045555/task_20250413_192014/conversation.json diff --git a/ta-browser/1744571356739037821/task_20250413_190916/conversation.json b/ta-browser/1744571356739037821/task_20250413_190916/conversation.json deleted file mode 100644 index e9976ca..0000000 --- a/ta-browser/1744571356739037821/task_20250413_190916/conversation.json +++ /dev/null @@ -1,79 +0,0 @@ -[ - { - "role": "user", - "content": "Inputs :\nUser Query : Go to Amazon.com and search for RTX 4060 graphics cards. Return the search results including prices and available models.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" - }, - { - "role": "assistant", - "content": "Plan: 1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\nNext Step: Navigate to Amazon's website (https://www.amazon.com)", - "name": "planner_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" current_step=\"Navigate to Amazon's website (https://www.amazon.com)\" " - }, - { - "role": "assistant", - "content": "I'll help navigate to Amazon's website using the open_url_tool.", - "name": "browser_nav_agent" - }, - { - "role": "assistant", - "content": null, - "tool_calls": [ - { - "id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", - "type": "function", - "function": { - "name": "open_url_tool", - "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" - } - } - ] - }, - { - "role": "tool", - "tool_call_id": "toolu_017ZWyYP2FhUsuRfeB5jc5TG", - "name": "open_url_tool", - "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" - }, - { - "role": "assistant", - "content": "I attempted to navigate to Amazon's website but encountered an error. This is common with Amazon as they have strict bot detection measures. Let me verify if we landed on any page by checking the current URL.", - "name": "browser_nav_agent" - }, - { - "role": "assistant", - "content": null, - "tool_calls": [ - { - "id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", - "type": "function", - "function": { - "name": "get_url", - "arguments": "{\"raw_args\": \"ArgsDict(args_dict={})\"}" - } - } - ] - }, - { - "role": "tool", - "tool_call_id": "toolu_01JL6GZfYgF6MpziTAG4mwBw", - "name": "get_url", - "content": "Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more." - }, - { - "role": "assistant", - "content": "Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.", - "name": "browser_nav_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\" next_step=\"Navigate to Amazon's website (https://www.amazon.com)\" tool_response=\"Despite the initial error message, we have successfully landed on Amazon's homepage. The URL and page title confirm we are on Amazon.com. The navigation step has been completed successfully, and we can proceed with the next step in the plan.\n\nSummary: Successfully navigated to Amazon.com website, verified by the current URL and page title. Ready for the next step to locate the search bar.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\nTool Call: get_url\nArguments: {}\nResponse: Current Page: https://www.amazon.com/, Title: Amazon.com. Spend less. Smile more.\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" - }, - { - "role": "assistant", - "content": "Feedback: Original Plan:\n1. Navigate to Amazon's website (https://www.amazon.com)\n2. Locate the search bar on Amazon's homepage\n3. Enter \"RTX 4060 graphics card\" in the search bar\n4. Submit the search query\n5. Wait for search results to load\n6. Check if results contain RTX 4060 graphics cards\n7. Extract product names, prices, and model information from the search results\n8. Compile the information about available models and their prices\n\nCurrent Progress:\nCurrently on Step 1: Navigate to Amazon's website (https://www.amazon.com)\n\nAnalysis and Feedback:\nThe browser agent has successfully completed the first step of navigating to Amazon.com. This is confirmed by:\n1. The get_url tool response showing we are at \"https://www.amazon.com/\"\n2. The page title verification showing \"Amazon.com. Spend less. Smile more.\"\n\nDespite an initial execution context error (which is common during page loads), the navigation was ultimately successful. The browser agent can now proceed to the next step of locating the search bar on Amazon's homepage.\n\nThe plan should continue with Step 2 as we have successfully completed Step 1.\nTerminate: False\nFinal Response: None", - "name": "critique_agent" - } -] \ No newline at end of file diff --git a/ta-browser/1744571514482277604/task_20250413_191154/conversation.json b/ta-browser/1744571514482277604/task_20250413_191154/conversation.json deleted file mode 100644 index 0637a08..0000000 --- a/ta-browser/1744571514482277604/task_20250413_191154/conversation.json +++ /dev/null @@ -1 +0,0 @@ -[] \ No newline at end of file diff --git a/ta-browser/1744572014568045555/task_20250413_192014/conversation.json b/ta-browser/1744572014568045555/task_20250413_192014/conversation.json deleted file mode 100644 index d214df6..0000000 --- a/ta-browser/1744572014568045555/task_20250413_192014/conversation.json +++ /dev/null @@ -1,49 +0,0 @@ -[ - { - "role": "user", - "content": "Inputs :\nUser Query : Go to Amazon website and search for \"rtx 4060\". Extract the product listings including prices, models, and ratings for RTX 4060 graphics cards.\nFeedback : None\nCurrent URL : https://www.google.com/\nCurrent Date : April 13, 2025" - }, - { - "role": "assistant", - "content": "Plan: 1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\nNext Step: Open Amazon's website via direct URL: https://www.amazon.com", - "name": "planner_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" current_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" " - }, - { - "role": "assistant", - "content": null, - "tool_calls": [ - { - "id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", - "type": "function", - "function": { - "name": "open_url_tool", - "arguments": "{\"url\": \"https://www.amazon.com\", \"timeout\": 3}" - } - } - ] - }, - { - "role": "tool", - "tool_call_id": "call_TUYaiyiEaKqnoTMlcZgpJqzM", - "name": "open_url_tool", - "content": "Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation" - }, - { - "role": "assistant", - "content": "I encountered an error while trying to open Amazon's website. The page failed to load successfully.", - "name": "browser_nav_agent" - }, - { - "role": "user", - "content": "Inputs :\nplan=\"1. Open Amazon's website via direct URL: https://www.amazon.com,\n2. Use the search bar to input 'RTX 4060',\n3. Submit the search query,\n4. Verify search results contain RTX 4060 listings,\n5. Extract product listings including prices, models, and ratings for RTX 4060 graphics cards,\n6. Compile extracted information for user.\" next_step=\"Open Amazon's website via direct URL: https://www.amazon.com\" tool_response=\"I encountered an error while trying to open Amazon's website. The page failed to load successfully.\" tool_interactions=\"Tool Call: open_url_tool\nArguments: {'url': 'https://www.amazon.com', 'timeout': 3}\nResponse: Error loading page: https://www.amazon.com, Error: Page.title: Execution context was destroyed, most likely because of a navigation\n---\n\" ss_analysis=\"SS analysis not available\"browser_error=\"None\"" - }, - { - "role": "assistant", - "content": "Feedback: The original plan involves opening Amazon's website and then performing a series of actions to search for 'RTX 4060' and extract product details. The current step was to open Amazon's website using the direct URL. However, the browser agent encountered an error while trying to load the page. The error message indicates that the execution context was destroyed, likely due to a navigation issue. This means that the page could not be successfully loaded, and thus we cannot proceed with the search and extraction tasks. Since this is a critical failure at the very first step, it's not possible to continue with the plan.\nTerminate: True\nFinal Response: The task encountered an error while trying to open Amazon's website, which prevented further actions. The error was related to the execution context being destroyed, likely due to a navigation issue. This means the page could not be loaded successfully, and as a result, we cannot proceed with searching for RTX 4060 graphics cards or extracting any product information. I recommend retrying the task or checking the network connection and URL for any issues.", - "name": "critique_agent" - } -] \ No newline at end of file diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 6a2aaef..69df1d7 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -261,7 +261,10 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com", m raise ValueError("job_id is required") # 1. Create conversation_storage with job_id - self.conversation_storage = ConversationStorage(job_id) + try: + self.conversation_storage = ConversationStorage(job_id=self.job_id) + except Exception as storage_error: + raise RuntimeError(f"Failed to initialize conversation storage: {str(storage_error)}") from storage_error # 2. Initialize browser manager try: diff --git a/ta-browser/core/utils/openai_msg_parser.py b/ta-browser/core/utils/openai_msg_parser.py index af57cdc..1458125 100644 --- a/ta-browser/core/utils/openai_msg_parser.py +++ b/ta-browser/core/utils/openai_msg_parser.py @@ -646,7 +646,6 @@ def get_conversation_history(self) -> List[ChatCompletionMessageParam]: return self.conversation_history -# In openai_msg_parser.py (or wherever ConversationStorage is defined) from datetime import datetime import os import json From c535c7742c5acbd2d4e0b2152df0c0ab4052938d Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Wed, 16 Apr 2025 15:52:56 +0530 Subject: [PATCH 15/24] feat: agent reinitialization everytime the model_preference changes --- cortex_on/agents/code_agent.py | 23 +++++---- cortex_on/agents/hello.py | 65 ++++++++++++++++++++++++++ cortex_on/agents/orchestrator_agent.py | 54 +++++++++++++++++---- cortex_on/agents/planner_agent.py | 29 +++++++----- cortex_on/instructor.py | 16 +++++-- cortex_on/main.py | 32 +++++++++++-- cortex_on/utils/ant_client.py | 24 ++++++++++ ta-browser/core/server/routes/web.py | 14 +++++- 8 files changed, 221 insertions(+), 36 deletions(-) create mode 100644 cortex_on/agents/hello.py diff --git a/cortex_on/agents/code_agent.py b/cortex_on/agents/code_agent.py index 43fa047..8580ab6 100644 --- a/cortex_on/agents/code_agent.py +++ b/cortex_on/agents/code_agent.py @@ -15,7 +15,7 @@ from pydantic_ai.models.anthropic import AnthropicModel # Local application imports -from utils.ant_client import get_client +from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client from utils.stream_response_format import StreamResponse load_dotenv() @@ -25,6 +25,7 @@ class CoderAgentDeps: websocket: Optional[WebSocket] = None stream_output: Optional[StreamResponse] = None + model_preference: str = "Anthropic" # Constants ALLOWED_COMMANDS = { @@ -236,20 +237,26 @@ async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> N stream_output_json = json.dumps(asdict(ctx.deps.stream_output)) logfire.debug("WebSocket message sent: {stream_output_json}", stream_output_json=stream_output_json) -# Initialize the model -model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() -) + # Initialize the agent -coder_agent = Agent( + +def coder_agent(model_preference: str = "Anthropic"): + if model_preference == "Anthropic": + model = get_anthropic_model_instance() + elif model_preference == "OpenAI": + model = get_openai_model_instance() + else: + raise ValueError(f"Unknown model_preference: {model_preference}") + print(f"[CODER_INIT] Creating coder agent with model: {model}") + coder_agent = Agent( model=model, name="Coder Agent", result_type=CoderResult, deps_type=CoderAgentDeps, system_prompt=coder_system_message -) + ) + return coder_agent @coder_agent.tool async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: diff --git a/cortex_on/agents/hello.py b/cortex_on/agents/hello.py new file mode 100644 index 0000000..2c453cc --- /dev/null +++ b/cortex_on/agents/hello.py @@ -0,0 +1,65 @@ + +@dataclass +class orchestrator_deps: + websocket: Optional[WebSocket] = None + stream_output: Optional[StreamResponse] = None + # Add a collection to track agent-specific streams + agent_responses: Optional[List[StreamResponse]] = None + model_preference: str = "Anthropic" + +async def get_model_preference() -> str: + """Fetch model preference from ta_browser API""" + try: + async with aiohttp.ClientSession() as session: + async with session.get("http://localhost:8000/api/v1/model/preference") as response: + if response.status == 200: + data = await response.json() + return data.get("model_preference", "Anthropic") + else: + logfire.error(f"Failed to get model preference: {response.status}") + return "Anthropic" + except Exception as e: + logfire.error(f"Error getting model preference: {str(e)}") + return "Anthropic" + +async def initialize_orchestrator_agent(): + """Initialize the orchestrator agent with model preference from ta_browser""" + model_preference = await get_model_preference() + model = AnthropicModel( + model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), + anthropic_client=get_client() + ) + + print(f"Orchestrator agent model initialized with MODEL PREFERENCE: {model_preference}") + + return Agent( + model=model, + name="Orchestrator Agent", + system_prompt=orchestrator_system_prompt, + deps_type=orchestrator_deps + ) + +# Initialize the agent +orchestrator_agent = None + +def get_orchestrator_agent(): + """Get the orchestrator agent, initializing it if necessary""" + global orchestrator_agent + if orchestrator_agent is None: + raise RuntimeError("Orchestrator agent not initialized. Call initialize_orchestrator_agent() first.") + return orchestrator_agent + +@get_orchestrator_agent().tool +async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Plans the task and assigns it to the appropriate agents""" + try: + logfire.info(f"Planning task: {task}") + + # Create a new StreamResponse for Planner Agent + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 284d927..fb92b3c 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -14,7 +14,9 @@ from utils.stream_response_format import StreamResponse from agents.planner_agent import planner_agent, update_todo_status from agents.code_agent import coder_agent, CoderAgentDeps -from utils.ant_client import get_client +from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client + +print("[ORCH_INIT] Starting orchestrator agent initialization") @dataclass class orchestrator_deps: @@ -24,6 +26,7 @@ class orchestrator_deps: agent_responses: Optional[List[StreamResponse]] = None model_preference: str = "Anthropic" + orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. [AGENT CAPABILITIES] @@ -159,17 +162,39 @@ class orchestrator_deps: - Format: "Task description (agent_name)" """ -model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() -) +# print("[ORCH_INIT] Getting Anthropic client...") +# client = get_client() +# print(f"[ORCH_INIT] Anthropic client obtained: {client}") + +print("[ORCH_INIT] Initializing Anthropic model...") +model = get_anthropic_model_instance() +print(f"[ORCH_INIT] Anthropic model initialized: {model}") + +print("Orchestrator agent model initialized without MODEL PREFERENCE") + +print("[ORCH_INIT] Creating orchestrator agent...") orchestrator_agent = Agent( model=model, name="Orchestrator Agent", system_prompt=orchestrator_system_prompt, deps_type=orchestrator_deps ) +print("[ORCH_INIT] Orchestrator agent created successfully") +def orchestrator_agent(model_preference: str): + if model_preference == "Anthropic": + model = get_anthropic_model_instance() + elif model_preference == "OpenAI": + model = get_openai_model_instance() + else: + raise ValueError(f"Unknown model_preference: {model_preference}") + print(f"[ORCH_INIT] Creating orchestrator agent with model: {model}") + return Agent( + model=model, + name="Orchestrator Agent", + system_prompt=orchestrator_system_prompt, + deps_type=orchestrator_deps + ) @orchestrator_agent.tool async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: @@ -197,7 +222,10 @@ async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) # Run planner agent - planner_response = await planner_agent.run(user_prompt=task) + agent = planner_agent(model_preference=ctx.deps.model_preference) + planner_response = await agent.run(user_prompt=task) + + logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") # Update planner stream with results plan_text = planner_response.data.plan @@ -253,14 +281,18 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: # Create deps with the new stream_output deps_for_coder_agent = CoderAgentDeps( websocket=ctx.deps.websocket, - stream_output=coder_stream_output + stream_output=coder_stream_output, + model_preference = ctx.deps.model_preference ) # Run coder agent - coder_response = await coder_agent.run( + agent = coder_agent(model_preference=ctx.deps.model_preference) + coder_response = await agent.run( user_prompt=task, deps=deps_for_coder_agent ) + logfire.info(f"coder_response: {coder_response}") + logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") # Extract response data response_data = coder_response.data.content @@ -313,6 +345,7 @@ async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: api_url="http://localhost:8000/api/v1/web/stream", model_preference=ctx.deps.model_preference ) + logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") # Run WebSurfer with its own stream_output success, message, messages = await web_surfer_agent.generate_reply( @@ -442,7 +475,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas # We'll directly call planner_agent.run() to create a new plan first plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - plan_response = await planner_agent.run(user_prompt=plan_prompt) + agent = planner_agent(model_preference=ctx.deps.model_preference) + plan_response = await agent.run(user_prompt=plan_prompt) current_content = plan_response.data.plan else: # Read existing todo.md @@ -488,7 +522,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.error(error_msg, exc_info=True) planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = a500 + planner_stream_output.status_code = 500 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) return f"Failed to update the plan: {error_msg}" diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 897c22f..53592ae 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -10,7 +10,7 @@ from pydantic_ai.models.anthropic import AnthropicModel # Local application imports -from utils.ant_client import get_client +from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client @@ -176,17 +176,24 @@ class PlannerResult(BaseModel): plan: str = Field(description="The generated or updated plan in string format - this should be the complete plan text") -model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() -) -planner_agent = Agent( - model=model, - name="Planner Agent", - result_type=PlannerResult, - system_prompt=planner_prompt -) + + +def planner_agent(model_preference: str = "Anthropic"): + if model_preference == "Anthropic": + model = get_anthropic_model_instance() + elif model_preference == "OpenAI": + model = get_openai_model_instance() + else: + raise ValueError(f"Unknown model_preference: {model_preference}") + print(f"[PLANNER_INIT] Creating planner agent with model: {model}") + planner_agent = Agent( + model=model, + name="Planner Agent", + result_type=PlannerResult, + system_prompt=planner_prompt + ) + return planner_agent @planner_agent.tool_plain async def update_todo_status(task_description: str) -> str: diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index f2688b4..d92cf4e 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -25,8 +25,10 @@ load_dotenv() - - +print("[INIT] Loading agents...") +print(f"[INIT] Planner agent loaded: {planner_agent}") +print(f"[INIT] Coder agent loaded: {coder_agent}") +print(f"[INIT] Orchestrator agent loaded: {get_orchestrator_agent}") class DateTimeEncoder(json.JSONEncoder): """Custom JSON encoder that can handle datetime objects""" @@ -39,11 +41,13 @@ def default(self, obj): # Main Orchestrator Class class SystemInstructor: def __init__(self, model_preference: str = "Anthropic"): + print(f"[INIT] Initializing SystemInstructor with model_preference: {model_preference}") self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None self.orchestrator_response: List[StreamResponse] = [] self.model_preference = model_preference self._setup_logging() + print("[INIT] SystemInstructor initialization complete") def _setup_logging(self) -> None: """Configure logging with proper formatting""" @@ -91,10 +95,16 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: stream_output.steps.append("Agents initialized successfully") await self._safe_websocket_send(stream_output) - orchestrator_response = await orchestrator_agent.run( + print("Orchestrator agent running") + + agent = orchestrator_agent(self.model_preference) + orchestrator_response = await agent.run( user_prompt=task, deps=deps_for_orchestrator ) + + logfire.info(f"Orchestrator Agent using model type: {self.model_preference}") + stream_output.output = orchestrator_response.data stream_output.status_code = 200 logfire.debug(f"Orchestrator response: {orchestrator_response.data}") diff --git a/cortex_on/main.py b/cortex_on/main.py index e1adddd..4bab62f 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -13,12 +13,25 @@ # Default model preference is Anthropic MODEL_PREFERENCE = "Anthropic" +# Global instructor instance +instructor = None + @asynccontextmanager async def lifespan(app: FastAPI): # Set default model preference at startup app.state.model_preference = MODEL_PREFERENCE print(f"[STARTUP] Setting default model preference to: {MODEL_PREFERENCE}") + + # Initialize the instructor + global instructor + instructor = SystemInstructor(model_preference=MODEL_PREFERENCE) + print("[STARTUP] Instructor initialized") + yield + + # Cleanup + if instructor: + await instructor.shutdown() app: FastAPI = FastAPI(lifespan=lifespan) @@ -42,21 +55,34 @@ async def get_model_preference() -> str: @app.get("/set_model_preference") async def set_model_preference(model: str): """ - Set the model preference (Anthropic or OpenAI) + Set the model preference (Anthropic or OpenAI) and reinitialize the instructor """ if model not in ["Anthropic", "OpenAI"]: print(f"[ERROR] Invalid model preference attempted: {model}") return {"error": "Model must be 'Anthropic' or 'OpenAI'"} + print(f"[SET] Changing model preference from {app.state.model_preference} to {model}") app.state.model_preference = model + + # Reinitialize the instructor with new model preference + global instructor + if instructor: + await instructor.shutdown() + instructor = SystemInstructor(model_preference=model) + print(f"[SET] Instructor reinitialized with model preference: {model}") + return {"message": f"Model preference set to {model}"} async def generate_response(task: str, websocket: Optional[WebSocket] = None, model_preference: str = None): if model_preference is None: model_preference = app.state.model_preference print(f"[GENERATE] Using model preference: {model_preference} for task: {task[:30]}...") - orchestrator: SystemInstructor = SystemInstructor(model_preference) - return await orchestrator.run(task, websocket) + + global instructor + if not instructor: + instructor = SystemInstructor(model_preference=model_preference) + + return await instructor.run(task, websocket) @app.get("/agent/chat") async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List: diff --git a/cortex_on/utils/ant_client.py b/cortex_on/utils/ant_client.py index 924f0c6..c04732c 100644 --- a/cortex_on/utils/ant_client.py +++ b/cortex_on/utils/ant_client.py @@ -1,5 +1,6 @@ from pydantic_ai.models.anthropic import AnthropicModel from anthropic import AsyncAnthropic +from openai import AsyncOpenAI import os from dotenv import load_dotenv @@ -12,3 +13,26 @@ def get_client(): max_retries=3, timeout=10000) return client + + + +def get_openai_client(): + api_key = os.getenv("OPENAI_API_KEY") + client = AsyncOpenAI(api_key=api_key, + max_retries=3, + timeout=10000) + return client + +def get_openai_model_instance(): + model_name = os.getenv("OPENAI_MODEL_NAME") + # Create model instance + from pydantic_ai.models.openai import OpenAIModel + model_instance = OpenAIModel(model_name=model_name, openai_client=get_openai_client()) + return model_instance + +def get_anthropic_model_instance(): + model_name = os.getenv("ANTHROPIC_MODEL_NAME") + # Create model instance + from pydantic_ai.models.anthropic import AnthropicModel + model_instance = AnthropicModel(model_name=model_name, anthropic_client=get_client()) + return model_instance diff --git a/ta-browser/core/server/routes/web.py b/ta-browser/core/server/routes/web.py index 0f0ee09..8dabe34 100644 --- a/ta-browser/core/server/routes/web.py +++ b/ta-browser/core/server/routes/web.py @@ -1,17 +1,19 @@ import asyncio -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, WebSocket from fastapi.responses import StreamingResponse from typing import AsyncGenerator import json import time from datetime import datetime from queue import Empty +import os from core.server.models.web import StreamRequestModel, StreamResponseModel from core.server.constants import GLOBAL_TIMEOUT from core.server.utils.timeout import timeout from core.server.utils.session_tracker import SessionTracker from core.utils.logger import Logger +from core.utils.init_client import initialize_client logger = Logger() @@ -150,4 +152,14 @@ async def stream_session( except Exception as e: logger.error(f"Error in /stream: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@router.get("/api/v1/model/preference") +async def get_model_preference(): + """Get the current model preference""" + try: + # Get model preference from environment or use default + model_preference = os.getenv("MODEL_PREFERENCE", "Anthropic") + return {"model_preference": model_preference} + except Exception as e: raise HTTPException(status_code=500, detail=str(e)) \ No newline at end of file From d69a17f9091ea3263cd591666b2722460ad623b4 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Wed, 16 Apr 2025 17:25:23 +0530 Subject: [PATCH 16/24] feat: agent reinitialization everytime the model_preference changes --- cortex_on/agents/code_agent.py | 296 +++++------ cortex_on/agents/orchestrator_agent.py | 670 +++++++++++++------------ cortex_on/agents/planner_agent.py | 296 +++++------ cortex_on/main.py | 6 +- 4 files changed, 641 insertions(+), 627 deletions(-) diff --git a/cortex_on/agents/code_agent.py b/cortex_on/agents/code_agent.py index 8580ab6..d1eb785 100644 --- a/cortex_on/agents/code_agent.py +++ b/cortex_on/agents/code_agent.py @@ -241,7 +241,7 @@ async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> N # Initialize the agent -def coder_agent(model_preference: str = "Anthropic"): +async def coder_agent(model_preference: str = "Anthropic") -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -256,167 +256,169 @@ def coder_agent(model_preference: str = "Anthropic"): deps_type=CoderAgentDeps, system_prompt=coder_system_message ) - return coder_agent - -@coder_agent.tool -async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: - """ - Executes a shell command within a restricted directory and returns the output. - This consolidated tool handles terminal commands and file operations. - """ - try: - # Extract base command for security checks and messaging - base_command = command.split()[0] if command.split() else "" - - # Send operation description message - operation_message = get_high_level_operation_message(command, base_command) - await send_stream_update(ctx, operation_message) - - logfire.info("Executing shell command: {command}", command=command) - - # Setup restricted directory - base_dir = os.path.abspath(os.path.dirname(__file__)) - restricted_dir = os.path.join(base_dir, "code_files") - os.makedirs(restricted_dir, exist_ok=True) - - # Security check - if base_command not in ALLOWED_COMMANDS: - await send_stream_update(ctx, "Operation not permitted") - return f"Error: Command '{base_command}' is not allowed for security reasons." - - # Change to restricted directory for execution - original_dir = os.getcwd() - os.chdir(restricted_dir) - + + @coder_agent.tool + async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: + """ + Executes a shell command within a restricted directory and returns the output. + This consolidated tool handles terminal commands and file operations. + """ try: - # Handle echo with redirection (file writing) - if ">" in command and base_command == "echo": - file_path = command.split(">", 1)[1].strip() - await send_stream_update(ctx, f"Writing content to {file_path}") - - # Parse command parts - parts = command.split(">", 1) - echo_cmd = parts[0].strip() - - # Extract content, removing quotes if present - content = echo_cmd[5:].strip() - if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): - content = content[1:-1] - - try: - with open(file_path, "w") as file: - file.write(content) - - await send_stream_update(ctx, f"File {file_path} created successfully") - return f"Successfully wrote to {file_path}" - except Exception as e: - error_msg = f"Error writing to file: {str(e)}" - await send_stream_update(ctx, f"Failed to create file {file_path}") - logfire.error(error_msg, exc_info=True) - return error_msg + # Extract base command for security checks and messaging + base_command = command.split()[0] if command.split() else "" - # Handle cat with here-document for multiline file writing - elif "<<" in command and base_command == "cat": - cmd_parts = command.split("<<", 1) - cat_part = cmd_parts[0].strip() - - # Extract filename for status message if possible - file_path = None - if ">" in cat_part: - file_path = cat_part.split(">", 1)[1].strip() - await send_stream_update(ctx, f"Creating file {file_path}") + # Send operation description message + operation_message = get_high_level_operation_message(command, base_command) + await send_stream_update(ctx, operation_message) + + logfire.info("Executing shell command: {command}", command=command) + + # Setup restricted directory + base_dir = os.path.abspath(os.path.dirname(__file__)) + restricted_dir = os.path.join(base_dir, "code_files") + os.makedirs(restricted_dir, exist_ok=True) + + # Security check + if base_command not in ALLOWED_COMMANDS: + await send_stream_update(ctx, "Operation not permitted") + return f"Error: Command '{base_command}' is not allowed for security reasons." + + # Change to restricted directory for execution + original_dir = os.getcwd() + os.chdir(restricted_dir) + + try: + # Handle echo with redirection (file writing) + if ">" in command and base_command == "echo": + file_path = command.split(">", 1)[1].strip() + await send_stream_update(ctx, f"Writing content to {file_path}") + + # Parse command parts + parts = command.split(">", 1) + echo_cmd = parts[0].strip() + + # Extract content, removing quotes if present + content = echo_cmd[5:].strip() + if (content.startswith('"') and content.endswith('"')) or \ + (content.startswith("'") and content.endswith("'")): + content = content[1:-1] + + try: + with open(file_path, "w") as file: + file.write(content) + + await send_stream_update(ctx, f"File {file_path} created successfully") + return f"Successfully wrote to {file_path}" + except Exception as e: + error_msg = f"Error writing to file: {str(e)}" + await send_stream_update(ctx, f"Failed to create file {file_path}") + logfire.error(error_msg, exc_info=True) + return error_msg - try: - # Parse heredoc parts - doc_part = cmd_parts[1].strip() + # Handle cat with here-document for multiline file writing + elif "<<" in command and base_command == "cat": + cmd_parts = command.split("<<", 1) + cat_part = cmd_parts[0].strip() - # Extract filename + # Extract filename for status message if possible + file_path = None if ">" in cat_part: file_path = cat_part.split(">", 1)[1].strip() - else: - await send_stream_update(ctx, "Invalid file operation") - return "Error: Invalid cat command format. Must include redirection." + await send_stream_update(ctx, f"Creating file {file_path}") - # Parse the heredoc content and delimiter - if "\n" in doc_part: - delimiter_and_content = doc_part.split("\n", 1) - delimiter = delimiter_and_content[0].strip("'").strip('"') - content = delimiter_and_content[1] + try: + # Parse heredoc parts + doc_part = cmd_parts[1].strip() - # Find the end delimiter and extract content - if f"\n{delimiter}" in content: - content = content.split(f"\n{delimiter}")[0] - - # Write to file - with open(file_path, "w") as file: - file.write(content) + # Extract filename + if ">" in cat_part: + file_path = cat_part.split(">", 1)[1].strip() + else: + await send_stream_update(ctx, "Invalid file operation") + return "Error: Invalid cat command format. Must include redirection." + + # Parse the heredoc content and delimiter + if "\n" in doc_part: + delimiter_and_content = doc_part.split("\n", 1) + delimiter = delimiter_and_content[0].strip("'").strip('"') + content = delimiter_and_content[1] - await send_stream_update(ctx, f"File {file_path} created successfully") - return f"Successfully wrote multiline content to {file_path}" + # Find the end delimiter and extract content + if f"\n{delimiter}" in content: + content = content.split(f"\n{delimiter}")[0] + + # Write to file + with open(file_path, "w") as file: + file.write(content) + + await send_stream_update(ctx, f"File {file_path} created successfully") + return f"Successfully wrote multiline content to {file_path}" + else: + await send_stream_update(ctx, "File content format error") + return "Error: End delimiter not found in heredoc" else: await send_stream_update(ctx, "File content format error") - return "Error: End delimiter not found in heredoc" - else: - await send_stream_update(ctx, "File content format error") - return "Error: Invalid heredoc format" - except Exception as e: - error_msg = f"Error processing cat with heredoc: {str(e)}" - file_path_str = file_path if file_path else 'file' - await send_stream_update(ctx, f"Failed to create file {file_path_str}") - logfire.error(error_msg, exc_info=True) - return error_msg - - # Execute standard commands - else: - # Send execution message - execution_msg = get_high_level_execution_message(command, base_command) - await send_stream_update(ctx, execution_msg) + return "Error: Invalid heredoc format" + except Exception as e: + error_msg = f"Error processing cat with heredoc: {str(e)}" + file_path_str = file_path if file_path else 'file' + await send_stream_update(ctx, f"Failed to create file {file_path_str}") + logfire.error(error_msg, exc_info=True) + return error_msg - # Execute the command using subprocess - try: - args = shlex.split(command) - result = subprocess.run( - args, - shell=True, - capture_output=True, - text=True, - timeout=60, - ) + # Execute standard commands + else: + # Send execution message + execution_msg = get_high_level_execution_message(command, base_command) + await send_stream_update(ctx, execution_msg) - logfire.info(f"Command executed: {result.args}") + # Execute the command using subprocess + try: + args = shlex.split(command) + result = subprocess.run( + args, + shell=True, + capture_output=True, + text=True, + timeout=60, + ) + + logfire.info(f"Command executed: {result.args}") + + # Handle success + if result.returncode == 0: + success_msg = get_success_message(command, base_command) + await send_stream_update(ctx, success_msg) + logfire.info(f"Command executed successfully: {result.stdout}") + return result.stdout + + # Handle failure + else: + files = os.listdir('.') + error_msg = f"Command failed with error code {result.returncode}:\n{result.stderr}\n\nFiles in directory: {files}" + failure_msg = get_failure_message(command, base_command) + await send_stream_update(ctx, failure_msg) + return error_msg - # Handle success - if result.returncode == 0: - success_msg = get_success_message(command, base_command) - await send_stream_update(ctx, success_msg) - logfire.info(f"Command executed successfully: {result.stdout}") - return result.stdout + except subprocess.TimeoutExpired: + await send_stream_update(ctx, "Operation timed out") + return "Command execution timed out after 60 seconds" - # Handle failure - else: - files = os.listdir('.') - error_msg = f"Command failed with error code {result.returncode}:\n{result.stderr}\n\nFiles in directory: {files}" - failure_msg = get_failure_message(command, base_command) - await send_stream_update(ctx, failure_msg) + except Exception as e: + error_msg = f"Error executing command: {str(e)}" + await send_stream_update(ctx, "Operation failed") + logfire.error(error_msg, exc_info=True) return error_msg + + finally: + # Always return to the original directory + os.chdir(original_dir) - except subprocess.TimeoutExpired: - await send_stream_update(ctx, "Operation timed out") - return "Command execution timed out after 60 seconds" - - except Exception as e: - error_msg = f"Error executing command: {str(e)}" - await send_stream_update(ctx, "Operation failed") - logfire.error(error_msg, exc_info=True) - return error_msg + except Exception as e: + error_msg = f"Error executing command: {str(e)}" + await send_stream_update(ctx, "Operation failed") + logfire.error(error_msg, exc_info=True) + return error_msg - finally: - # Always return to the original directory - os.chdir(original_dir) - - except Exception as e: - error_msg = f"Error executing command: {str(e)}" - await send_stream_update(ctx, "Operation failed") - logfire.error(error_msg, exc_info=True) - return error_msg \ No newline at end of file + logfire.info("All tools initialized for coder agent") + return coder_agent \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index fb92b3c..4a5eec5 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -166,22 +166,22 @@ class orchestrator_deps: # client = get_client() # print(f"[ORCH_INIT] Anthropic client obtained: {client}") -print("[ORCH_INIT] Initializing Anthropic model...") -model = get_anthropic_model_instance() - -print(f"[ORCH_INIT] Anthropic model initialized: {model}") - -print("Orchestrator agent model initialized without MODEL PREFERENCE") - -print("[ORCH_INIT] Creating orchestrator agent...") -orchestrator_agent = Agent( - model=model, - name="Orchestrator Agent", - system_prompt=orchestrator_system_prompt, - deps_type=orchestrator_deps -) -print("[ORCH_INIT] Orchestrator agent created successfully") -def orchestrator_agent(model_preference: str): +# print("[ORCH_INIT] Initializing Anthropic model...") +# model = get_anthropic_model_instance() + +# print(f"[ORCH_INIT] Anthropic model initialized: {model}") + +# print("Orchestrator agent model initialized without MODEL PREFERENCE") + +# print("[ORCH_INIT] Creating orchestrator agent...") +# orchestrator_agent = Agent( +# model=model, +# name="Orchestrator Agent", +# system_prompt=orchestrator_system_prompt, +# deps_type=orchestrator_deps +# ) +# print("[ORCH_INIT] Orchestrator agent created successfully") +async def orchestrator_agent(model_preference: str) -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -189,354 +189,364 @@ def orchestrator_agent(model_preference: str): else: raise ValueError(f"Unknown model_preference: {model_preference}") print(f"[ORCH_INIT] Creating orchestrator agent with model: {model}") - return Agent( + orchestrator_agent = Agent( model=model, name="Orchestrator Agent", system_prompt=orchestrator_system_prompt, deps_type=orchestrator_deps ) - -@orchestrator_agent.tool -async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Plans the task and assigns it to the appropriate agents""" - try: - logfire.info(f"Planning task: {task}") - # Create a new StreamResponse for Planner Agent - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(planner_stream_output) + @orchestrator_agent.tool + async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Plans the task and assigns it to the appropriate agents""" + try: + logfire.info(f"Planning task: {task}") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Update planner stream - planner_stream_output.steps.append("Planning task...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Run planner agent - agent = planner_agent(model_preference=ctx.deps.model_preference) - planner_response = await agent.run(user_prompt=task) + # Create a new StreamResponse for Planner Agent + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(planner_stream_output) + + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Update planner stream + planner_stream_output.steps.append("Planning task...") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Run planner agent + agent = planner_agent(model_preference=ctx.deps.model_preference) + planner_response = await agent.run(user_prompt=task) - logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") - - # Update planner stream with results - plan_text = planner_response.data.plan - planner_stream_output.steps.append("Task planned successfully") - planner_stream_output.output = plan_text - planner_stream_output.status_code = 200 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Also update orchestrator stream - ctx.deps.stream_output.steps.append("Task planned successfully") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - - return f"Task planned successfully\nTask: {plan_text}" - except Exception as e: - error_msg = f"Error planning task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update planner stream with error - if planner_stream_output: - planner_stream_output.steps.append(f"Planning failed: {str(e)}") - planner_stream_output.status_code = 500 + logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") + + # Update planner stream with results + plan_text = planner_response.data.plan + planner_stream_output.steps.append("Task planned successfully") + planner_stream_output.output = plan_text + planner_stream_output.status_code = 200 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # Also update orchestrator stream - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Planning failed: {str(e)}") + # Also update orchestrator stream + ctx.deps.stream_output.steps.append("Task planned successfully") await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - return f"Failed to plan task: {error_msg}" + return f"Task planned successfully\nTask: {plan_text}" + except Exception as e: + error_msg = f"Error planning task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update planner stream with error + if planner_stream_output: + planner_stream_output.steps.append(f"Planning failed: {str(e)}") + planner_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Also update orchestrator stream + if ctx.deps.stream_output: + ctx.deps.stream_output.steps.append(f"Planning failed: {str(e)}") + await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + + return f"Failed to plan task: {error_msg}" -@orchestrator_agent.tool -async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Assigns coding tasks to the coder agent""" - try: - logfire.info(f"Assigning coding task: {task}") - - # Create a new StreamResponse for Coder Agent - coder_stream_output = StreamResponse( - agent_name="Coder Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(coder_stream_output) - - # Send initial update for Coder Agent - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - # Create deps with the new stream_output - deps_for_coder_agent = CoderAgentDeps( - websocket=ctx.deps.websocket, - stream_output=coder_stream_output, - model_preference = ctx.deps.model_preference - ) - - # Run coder agent - agent = coder_agent(model_preference=ctx.deps.model_preference) - coder_response = await agent.run( - user_prompt=task, - deps=deps_for_coder_agent - ) - logfire.info(f"coder_response: {coder_response}") - logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") - - # Extract response data - response_data = coder_response.data.content - - # Update coder_stream_output with coding results - coder_stream_output.output = response_data - coder_stream_output.status_code = 200 - coder_stream_output.steps.append("Coding task completed successfully") - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - # Add a reminder in the result message to update the plan using planner_agent_update - response_with_reminder = f"{response_data}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" - - return response_with_reminder - except Exception as e: - error_msg = f"Error assigning coding task: {str(e)}" - logfire.error(error_msg, exc_info=True) + @orchestrator_agent.tool + async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Assigns coding tasks to the coder agent""" + try: + logfire.info(f"Assigning coding task: {task}") + + # Create a new StreamResponse for Coder Agent + coder_stream_output = StreamResponse( + agent_name="Coder Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(coder_stream_output) + + # Send initial update for Coder Agent + await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + + # Create deps with the new stream_output + deps_for_coder_agent = CoderAgentDeps( + websocket=ctx.deps.websocket, + stream_output=coder_stream_output, + model_preference = ctx.deps.model_preference + ) + + # Run coder agent + agent = coder_agent(model_preference=ctx.deps.model_preference) + coder_response = await agent.run( + user_prompt=task, + deps=deps_for_coder_agent + ) + logfire.info(f"coder_response: {coder_response}") + logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") + + # Extract response data + response_data = coder_response.data.content + + # Update coder_stream_output with coding results + coder_stream_output.output = response_data + coder_stream_output.status_code = 200 + coder_stream_output.steps.append("Coding task completed successfully") + await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + + # Add a reminder in the result message to update the plan using planner_agent_update + response_with_reminder = f"{response_data}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" + + return response_with_reminder + except Exception as e: + error_msg = f"Error assigning coding task: {str(e)}" + logfire.error(error_msg, exc_info=True) - # Update coder_stream_output with error - coder_stream_output.steps.append(f"Coding task failed: {str(e)}") - coder_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + # Update coder_stream_output with error + coder_stream_output.steps.append(f"Coding task failed: {str(e)}") + coder_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - return f"Failed to assign coding task: {error_msg}" + return f"Failed to assign coding task: {error_msg}" -@orchestrator_agent.tool -async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Assigns web surfing tasks to the web surfer agent""" - try: - logfire.info(f"Assigning web surfing task: {task}") - - # Create a new StreamResponse for WebSurfer - web_surfer_stream_output = StreamResponse( - agent_name="Web Surfer", - instructions=task, - steps=[], - output="", - status_code=0, - live_url=None - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(web_surfer_stream_output) - - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - # Initialize WebSurfer agent with model preference - web_surfer_agent = WebSurfer( - api_url="http://localhost:8000/api/v1/web/stream", - model_preference=ctx.deps.model_preference - ) - logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") - - # Run WebSurfer with its own stream_output - success, message, messages = await web_surfer_agent.generate_reply( - instruction=task, - websocket=ctx.deps.websocket, - stream_output=web_surfer_stream_output - ) - - # Update WebSurfer's stream_output with final result - if success: - web_surfer_stream_output.steps.append("Web search completed successfully") - web_surfer_stream_output.output = message - web_surfer_stream_output.status_code = 200 - - # Add a reminder to update the plan - message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" - else: - web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") + @orchestrator_agent.tool + async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: + """Assigns web surfing tasks to the web surfer agent""" + try: + logfire.info(f"Assigning web surfing task: {task}") + + # Create a new StreamResponse for WebSurfer + web_surfer_stream_output = StreamResponse( + agent_name="Web Surfer", + instructions=task, + steps=[], + output="", + status_code=0, + live_url=None + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(web_surfer_stream_output) + + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + + # Initialize WebSurfer agent with model preference + web_surfer_agent = WebSurfer( + api_url="http://localhost:8000/api/v1/web/stream", + model_preference=ctx.deps.model_preference + ) + logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") + + # Run WebSurfer with its own stream_output + success, message, messages = await web_surfer_agent.generate_reply( + instruction=task, + websocket=ctx.deps.websocket, + stream_output=web_surfer_stream_output + ) + + # Update WebSurfer's stream_output with final result + if success: + web_surfer_stream_output.steps.append("Web search completed successfully") + web_surfer_stream_output.output = message + web_surfer_stream_output.status_code = 200 + + # Add a reminder to update the plan + message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" + else: + web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") + web_surfer_stream_output.status_code = 500 + message_with_reminder = message + + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + + web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + + return message_with_reminder + except Exception as e: + error_msg = f"Error assigning web surfing task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update WebSurfer's stream_output with error + web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") web_surfer_stream_output.status_code = 500 - message_with_reminder = message - - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - return message_with_reminder - except Exception as e: - error_msg = f"Error assigning web surfing task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update WebSurfer's stream_output with error - web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") - web_surfer_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - return f"Failed to assign web surfing task: {error_msg}" - -@orchestrator_agent.tool -async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: - """Sends a question to the frontend and waits for human input""" - try: - logfire.info(f"Asking human: {question}") - - # Create a new StreamResponse for Human Input - human_stream_output = StreamResponse( - agent_name="Human Input", - instructions=question, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(human_stream_output) - - # Send the question to frontend - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - # Update stream with waiting message - human_stream_output.steps.append("Waiting for human input...") - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - # Wait for response from frontend - response = await ctx.deps.websocket.receive_text() - - # Update stream with response - human_stream_output.steps.append("Received human input") - human_stream_output.output = response - human_stream_output.status_code = 200 - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - return response - except Exception as e: - error_msg = f"Error getting human input: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update stream with error - human_stream_output.steps.append(f"Failed to get human input: {str(e)}") - human_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - - return f"Failed to get human input: {error_msg}" + await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + return f"Failed to assign web surfing task: {error_msg}" -@orchestrator_agent.tool -async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_task: str) -> str: - """ - Updates the todo.md file to mark a task as completed and returns the full updated plan. - - Args: - completed_task: Description of the completed task including which agent performed it - - Returns: - The complete updated todo.md content with tasks marked as completed - """ - try: - logfire.info(f"Updating plan with completed task: {completed_task}") - - # Create a new StreamResponse for Planner Agent update - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=f"Update todo.md to mark as completed: {completed_task}", - steps=[], - output="", - status_code=0 - ) - - # Send initial update - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Directly read and update the todo.md file - base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) - planner_dir = os.path.join(base_dir, "agents", "planner") - todo_path = os.path.join(planner_dir, "todo.md") - - planner_stream_output.steps.append("Reading current todo.md...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Make sure the directory exists - os.makedirs(planner_dir, exist_ok=True) - + @orchestrator_agent.tool + async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: + """Sends a question to the frontend and waits for human input""" try: - # Check if todo.md exists - if not os.path.exists(todo_path): - planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # We'll directly call planner_agent.run() to create a new plan first - plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - agent = planner_agent(model_preference=ctx.deps.model_preference) - plan_response = await agent.run(user_prompt=plan_prompt) - current_content = plan_response.data.plan - else: - # Read existing todo.md - with open(todo_path, "r") as file: - current_content = file.read() - planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + logfire.info(f"Asking human: {question}") - # Now call planner_agent.run() with specific instructions to update the plan - update_prompt = f""" - Here is the current todo.md content: + # Create a new StreamResponse for Human Input + human_stream_output = StreamResponse( + agent_name="Human Input", + instructions=question, + steps=[], + output="", + status_code=0 + ) + + # Add to orchestrator's response collection if available + if ctx.deps.agent_responses is not None: + ctx.deps.agent_responses.append(human_stream_output) + + # Send the question to frontend + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - {current_content} + # Update stream with waiting message + human_stream_output.steps.append("Waiting for human input...") + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - Please update this plan to mark the following task as completed: {completed_task} - Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. - """ + # Wait for response from frontend + response = await ctx.deps.websocket.receive_text() - planner_stream_output.steps.append("Asking planner to update the plan...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + # Update stream with response + human_stream_output.steps.append("Received human input") + human_stream_output.output = response + human_stream_output.status_code = 200 + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - updated_plan_response = await planner_agent.run(user_prompt=update_prompt) - updated_plan = updated_plan_response.data.plan + return response + except Exception as e: + error_msg = f"Error getting human input: {str(e)}" + logfire.error(error_msg, exc_info=True) - # Write the updated plan back to todo.md - with open(todo_path, "w") as file: - file.write(updated_plan) + # Update stream with error + human_stream_output.steps.append(f"Failed to get human input: {str(e)}") + human_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, human_stream_output) - planner_stream_output.steps.append("Plan updated successfully") - planner_stream_output.output = updated_plan - planner_stream_output.status_code = 200 + return f"Failed to get human input: {error_msg}" + + @orchestrator_agent.tool + async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_task: str) -> str: + """ + Updates the todo.md file to mark a task as completed and returns the full updated plan. + + Args: + completed_task: Description of the completed task including which agent performed it + + Returns: + The complete updated todo.md content with tasks marked as completed + """ + try: + logfire.info(f"Updating plan with completed task: {completed_task}") + + # Create a new StreamResponse for Planner Agent update + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=f"Update todo.md to mark as completed: {completed_task}", + steps=[], + output="", + status_code=0 + ) + + # Send initial update await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # Update orchestrator stream - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Plan updated to mark task as completed: {completed_task}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + # Directly read and update the todo.md file + base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) + planner_dir = os.path.join(base_dir, "agents", "planner") + todo_path = os.path.join(planner_dir, "todo.md") + + planner_stream_output.steps.append("Reading current todo.md...") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Make sure the directory exists + os.makedirs(planner_dir, exist_ok=True) - return updated_plan + try: + # Check if todo.md exists + if not os.path.exists(todo_path): + planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # We'll directly call planner_agent.run() to create a new plan first + plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" + agent = planner_agent(model_preference=ctx.deps.model_preference) + plan_response = await agent.run(user_prompt=plan_prompt) + current_content = plan_response.data.plan + + logfire.info(f"Updating plan task using model type: {ctx.deps.model_preference}") + + else: + # Read existing todo.md + with open(todo_path, "r") as file: + current_content = file.read() + planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Now call planner_agent.run() with specific instructions to update the plan + update_prompt = f""" + Here is the current todo.md content: + + {current_content} + + Please update this plan to mark the following task as completed: {completed_task} + Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. + """ + + planner_stream_output.steps.append("Asking planner to update the plan...") + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + agent = planner_agent(model_preference=ctx.deps.model_preference) + updated_plan_response = await agent.run(user_prompt=update_prompt) + updated_plan = updated_plan_response.data.plan + + logfire.info(f"Updating prompt for plan task using model type: {ctx.deps.model_preference}") + + # Write the updated plan back to todo.md + with open(todo_path, "w") as file: + file.write(updated_plan) + + planner_stream_output.steps.append("Plan updated successfully") + planner_stream_output.output = updated_plan + planner_stream_output.status_code = 200 + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + # Update orchestrator stream + if ctx.deps.stream_output: + ctx.deps.stream_output.steps.append(f"Plan updated to mark task as completed: {completed_task}") + await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + + return updated_plan + + except Exception as e: + error_msg = f"Error during plan update operations: {str(e)}" + logfire.error(error_msg, exc_info=True) + + planner_stream_output.steps.append(f"Plan update failed: {str(e)}") + planner_stream_output.status_code = 500 + await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + + return f"Failed to update the plan: {error_msg}" except Exception as e: - error_msg = f"Error during plan update operations: {str(e)}" + error_msg = f"Error updating plan: {str(e)}" logfire.error(error_msg, exc_info=True) - planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + # Update stream output with error + if ctx.deps.stream_output: + ctx.deps.stream_output.steps.append(f"Failed to update plan: {str(e)}") + await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - return f"Failed to update the plan: {error_msg}" - - except Exception as e: - error_msg = f"Error updating plan: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update stream output with error - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Failed to update plan: {str(e)}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - - return f"Failed to update plan: {error_msg}" + return f"Failed to update plan: {error_msg}" + + logfire.info("All tools initialized for orchestrator agent") + return orchestrator_agent + # Helper function for sending WebSocket messages async def _safe_websocket_send(websocket: Optional[WebSocket], message: Any) -> bool: diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 53592ae..5d29b7d 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -179,7 +179,7 @@ class PlannerResult(BaseModel): -def planner_agent(model_preference: str = "Anthropic"): +async def planner_agent(model_preference: str = "Anthropic") -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -193,162 +193,164 @@ def planner_agent(model_preference: str = "Anthropic"): result_type=PlannerResult, system_prompt=planner_prompt ) - return planner_agent -@planner_agent.tool_plain -async def update_todo_status(task_description: str) -> str: - """ - A helper function that logs the update request but lets the planner agent handle the actual update logic. - - Args: - task_description: Description of the completed task + @planner_agent.tool_plain + async def update_todo_status(task_description: str) -> str: + """ + A helper function that logs the update request but lets the planner agent handle the actual update logic. - Returns: - A simple log message - """ - logfire.info(f"Received request to update todo.md for task: {task_description}") - return f"Received update request for: {task_description}" + Args: + task_description: Description of the completed task + + Returns: + A simple log message + """ + logfire.info(f"Received request to update todo.md for task: {task_description}") + return f"Received update request for: {task_description}" -@planner_agent.tool_plain -async def execute_terminal(command: str) -> str: - """ - Executes a terminal command within the planner directory for file operations. - This consolidated tool handles reading and writing plan files. - Restricted to only read and write operations for security. - """ - try: - logfire.info(f"Executing terminal command: {command}") - - # Define the restricted directory - base_dir = os.path.abspath(os.path.dirname(__file__)) - planner_dir = os.path.join(base_dir, "planner") - os.makedirs(planner_dir, exist_ok=True) - - # Extract the base command - base_command = command.split()[0] - - # Allow only read and write operations - ALLOWED_COMMANDS = {"cat", "echo", "ls"} - - # Security checks - if base_command not in ALLOWED_COMMANDS: - return f"Error: Command '{base_command}' is not allowed. Only read and write operations are permitted." - - if ".." in command or "~" in command or "/" in command: - return "Error: Path traversal attempts are not allowed." - - # Change to the restricted directory - original_dir = os.getcwd() - os.chdir(planner_dir) - + @planner_agent.tool_plain + async def execute_terminal(command: str) -> str: + """ + Executes a terminal command within the planner directory for file operations. + This consolidated tool handles reading and writing plan files. + Restricted to only read and write operations for security. + """ try: - # Handle echo with >> (append) - if base_command == "echo" and ">>" in command: - try: - # Split only on the first occurrence of >> - parts = command.split(">>", 1) - echo_part = parts[0].strip() + logfire.info(f"Executing terminal command: {command}") + + # Define the restricted directory + base_dir = os.path.abspath(os.path.dirname(__file__)) + planner_dir = os.path.join(base_dir, "planner") + os.makedirs(planner_dir, exist_ok=True) + + # Extract the base command + base_command = command.split()[0] + + # Allow only read and write operations + ALLOWED_COMMANDS = {"cat", "echo", "ls"} + + # Security checks + if base_command not in ALLOWED_COMMANDS: + return f"Error: Command '{base_command}' is not allowed. Only read and write operations are permitted." + + if ".." in command or "~" in command or "/" in command: + return "Error: Path traversal attempts are not allowed." + + # Change to the restricted directory + original_dir = os.getcwd() + os.chdir(planner_dir) + + try: + # Handle echo with >> (append) + if base_command == "echo" and ">>" in command: + try: + # Split only on the first occurrence of >> + parts = command.split(">>", 1) + echo_part = parts[0].strip() + file_path = parts[1].strip() + + # Extract content after echo command + content = echo_part[4:].strip() + + # Handle quotes if present + if (content.startswith('"') and content.endswith('"')) or \ + (content.startswith("'") and content.endswith("'")): + content = content[1:-1] + + # Append to file + with open(file_path, "a") as file: + file.write(content + "\n") + return f"Successfully appended to {file_path}" + except Exception as e: + logfire.error(f"Error appending to file: {str(e)}", exc_info=True) + return f"Error appending to file: {str(e)}" + + # Special handling for echo with redirection (file writing) + elif ">" in command and base_command == "echo" and ">>" not in command: + # Simple parsing for echo "content" > file.txt + parts = command.split(">", 1) + echo_cmd = parts[0].strip() file_path = parts[1].strip() - # Extract content after echo command - content = echo_part[4:].strip() - - # Handle quotes if present + # Extract content between echo and > (removing quotes if present) + content = echo_cmd[5:].strip() if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): + (content.startswith("'") and content.endswith("'")): content = content[1:-1] - # Append to file - with open(file_path, "a") as file: - file.write(content + "\n") - return f"Successfully appended to {file_path}" - except Exception as e: - logfire.error(f"Error appending to file: {str(e)}", exc_info=True) - return f"Error appending to file: {str(e)}" - - # Special handling for echo with redirection (file writing) - elif ">" in command and base_command == "echo" and ">>" not in command: - # Simple parsing for echo "content" > file.txt - parts = command.split(">", 1) - echo_cmd = parts[0].strip() - file_path = parts[1].strip() + # Write to file + try: + with open(file_path, "w") as file: + file.write(content) + return f"Successfully wrote to {file_path}" + except Exception as e: + logfire.error(f"Error writing to file: {str(e)}", exc_info=True) + return f"Error writing to file: {str(e)}" - # Extract content between echo and > (removing quotes if present) - content = echo_cmd[5:].strip() - if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): - content = content[1:-1] - - # Write to file - try: - with open(file_path, "w") as file: - file.write(content) - return f"Successfully wrote to {file_path}" - except Exception as e: - logfire.error(f"Error writing to file: {str(e)}", exc_info=True) - return f"Error writing to file: {str(e)}" - - # Handle cat with here-document for multiline file writing - elif "<<" in command and base_command == "cat": - try: - # Parse the command: cat > file.md << 'EOF'\nplan content\nEOF - cmd_parts = command.split("<<", 1) - cat_part = cmd_parts[0].strip() - doc_part = cmd_parts[1].strip() - - # Extract filename - if ">" in cat_part: - file_path = cat_part.split(">", 1)[1].strip() - else: - return "Error: Invalid cat command format. Must include redirection." - - # Parse the heredoc content - if "\n" in doc_part: - delimiter_and_content = doc_part.split("\n", 1) - delimiter = delimiter_and_content[0].strip("'").strip('"') - content = delimiter_and_content[1] + # Handle cat with here-document for multiline file writing + elif "<<" in command and base_command == "cat": + try: + # Parse the command: cat > file.md << 'EOF'\nplan content\nEOF + cmd_parts = command.split("<<", 1) + cat_part = cmd_parts[0].strip() + doc_part = cmd_parts[1].strip() + + # Extract filename + if ">" in cat_part: + file_path = cat_part.split(">", 1)[1].strip() + else: + return "Error: Invalid cat command format. Must include redirection." - # Find the end delimiter and extract content - if f"\n{delimiter}" in content: - content = content.split(f"\n{delimiter}")[0] + # Parse the heredoc content + if "\n" in doc_part: + delimiter_and_content = doc_part.split("\n", 1) + delimiter = delimiter_and_content[0].strip("'").strip('"') + content = delimiter_and_content[1] - # Write to file - with open(file_path, "w") as file: - file.write(content) - return f"Successfully wrote multiline content to {file_path}" + # Find the end delimiter and extract content + if f"\n{delimiter}" in content: + content = content.split(f"\n{delimiter}")[0] + + # Write to file + with open(file_path, "w") as file: + file.write(content) + return f"Successfully wrote multiline content to {file_path}" + else: + return "Error: End delimiter not found in heredoc" else: - return "Error: End delimiter not found in heredoc" - else: - return "Error: Invalid heredoc format" - except Exception as e: - logfire.error(f"Error processing cat with heredoc: {str(e)}", exc_info=True) - return f"Error processing cat with heredoc: {str(e)}" - - # Handle cat for reading files - elif base_command == "cat" and ">" not in command and "<<" not in command: - try: - file_path = command.split()[1] - with open(file_path, "r") as file: - content = file.read() - return content - except Exception as e: - logfire.error(f"Error reading file: {str(e)}", exc_info=True) - return f"Error reading file: {str(e)}" - - # Handle ls for listing files - elif base_command == "ls": - try: - files = os.listdir('.') - return "Files in planner directory:\n" + "\n".join(files) - except Exception as e: - logfire.error(f"Error listing files: {str(e)}", exc_info=True) - return f"Error listing files: {str(e)}" - else: - return f"Error: Command '{command}' is not supported. Only read and write operations are permitted." - - finally: - os.chdir(original_dir) - - except Exception as e: - logfire.error(f"Error executing command: {str(e)}", exc_info=True) - return f"Error executing command: {str(e)}" \ No newline at end of file + return "Error: Invalid heredoc format" + except Exception as e: + logfire.error(f"Error processing cat with heredoc: {str(e)}", exc_info=True) + return f"Error processing cat with heredoc: {str(e)}" + + # Handle cat for reading files + elif base_command == "cat" and ">" not in command and "<<" not in command: + try: + file_path = command.split()[1] + with open(file_path, "r") as file: + content = file.read() + return content + except Exception as e: + logfire.error(f"Error reading file: {str(e)}", exc_info=True) + return f"Error reading file: {str(e)}" + + # Handle ls for listing files + elif base_command == "ls": + try: + files = os.listdir('.') + return "Files in planner directory:\n" + "\n".join(files) + except Exception as e: + logfire.error(f"Error listing files: {str(e)}", exc_info=True) + return f"Error listing files: {str(e)}" + else: + return f"Error: Command '{command}' is not supported. Only read and write operations are permitted." + + finally: + os.chdir(original_dir) + + except Exception as e: + logfire.error(f"Error executing command: {str(e)}", exc_info=True) + return f"Error executing command: {str(e)}" + + logfire.info("All tools initialized for planner agent") + return planner_agent \ No newline at end of file diff --git a/cortex_on/main.py b/cortex_on/main.py index 4bab62f..4488c3a 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -29,9 +29,9 @@ async def lifespan(app: FastAPI): yield - # Cleanup - if instructor: - await instructor.shutdown() + # # Cleanup + # if instructor: + # await instructor.shutdown() app: FastAPI = FastAPI(lifespan=lifespan) From af6ea8a98a275840fe6e1646ca6899c6df6af8ec Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Wed, 16 Apr 2025 22:59:52 +0530 Subject: [PATCH 17/24] feat: agent reinitialization everytime the model_preference changes --- cortex_on/agents/hello.py | 65 -------------------------- cortex_on/agents/orchestrator_agent.py | 16 ++++--- cortex_on/instructor.py | 4 +- ta-browser/core/server/routes/web.py | 9 ---- 4 files changed, 11 insertions(+), 83 deletions(-) delete mode 100644 cortex_on/agents/hello.py diff --git a/cortex_on/agents/hello.py b/cortex_on/agents/hello.py deleted file mode 100644 index 2c453cc..0000000 --- a/cortex_on/agents/hello.py +++ /dev/null @@ -1,65 +0,0 @@ - -@dataclass -class orchestrator_deps: - websocket: Optional[WebSocket] = None - stream_output: Optional[StreamResponse] = None - # Add a collection to track agent-specific streams - agent_responses: Optional[List[StreamResponse]] = None - model_preference: str = "Anthropic" - -async def get_model_preference() -> str: - """Fetch model preference from ta_browser API""" - try: - async with aiohttp.ClientSession() as session: - async with session.get("http://localhost:8000/api/v1/model/preference") as response: - if response.status == 200: - data = await response.json() - return data.get("model_preference", "Anthropic") - else: - logfire.error(f"Failed to get model preference: {response.status}") - return "Anthropic" - except Exception as e: - logfire.error(f"Error getting model preference: {str(e)}") - return "Anthropic" - -async def initialize_orchestrator_agent(): - """Initialize the orchestrator agent with model preference from ta_browser""" - model_preference = await get_model_preference() - model = AnthropicModel( - model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() - ) - - print(f"Orchestrator agent model initialized with MODEL PREFERENCE: {model_preference}") - - return Agent( - model=model, - name="Orchestrator Agent", - system_prompt=orchestrator_system_prompt, - deps_type=orchestrator_deps - ) - -# Initialize the agent -orchestrator_agent = None - -def get_orchestrator_agent(): - """Get the orchestrator agent, initializing it if necessary""" - global orchestrator_agent - if orchestrator_agent is None: - raise RuntimeError("Orchestrator agent not initialized. Call initialize_orchestrator_agent() first.") - return orchestrator_agent - -@get_orchestrator_agent().tool -async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Plans the task and assigns it to the appropriate agents""" - try: - logfire.info(f"Planning task: {task}") - - # Create a new StreamResponse for Planner Agent - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 4a5eec5..2bfda6a 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -12,7 +12,7 @@ from pydantic_ai import Agent, RunContext from agents.web_surfer import WebSurfer from utils.stream_response_format import StreamResponse -from agents.planner_agent import planner_agent, update_todo_status +from agents.planner_agent import planner_agent from agents.code_agent import coder_agent, CoderAgentDeps from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client @@ -181,7 +181,7 @@ class orchestrator_deps: # deps_type=orchestrator_deps # ) # print("[ORCH_INIT] Orchestrator agent created successfully") -async def orchestrator_agent(model_preference: str) -> Agent: +async def orchestrator_agent(model_preference: str) : if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": @@ -222,7 +222,7 @@ async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) # Run planner agent - agent = planner_agent(model_preference=ctx.deps.model_preference) + agent = await planner_agent(model_preference=ctx.deps.model_preference) planner_response = await agent.run(user_prompt=task) logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") @@ -286,7 +286,7 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: ) # Run coder agent - agent = coder_agent(model_preference=ctx.deps.model_preference) + agent = await coder_agent(model_preference=ctx.deps.model_preference) coder_response = await agent.run( user_prompt=task, deps=deps_for_coder_agent @@ -475,7 +475,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas # We'll directly call planner_agent.run() to create a new plan first plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - agent = planner_agent(model_preference=ctx.deps.model_preference) + + agent = await planner_agent(model_preference=ctx.deps.model_preference) plan_response = await agent.run(user_prompt=plan_prompt) current_content = plan_response.data.plan @@ -501,7 +502,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas planner_stream_output.steps.append("Asking planner to update the plan...") await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - agent = planner_agent(model_preference=ctx.deps.model_preference) + agent = await planner_agent(model_preference=ctx.deps.model_preference) updated_plan_response = await agent.run(user_prompt=update_prompt) updated_plan = updated_plan_response.data.plan @@ -544,7 +545,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas return f"Failed to update plan: {error_msg}" - logfire.info("All tools initialized for orchestrator agent") + logfire.info("All tools initialized for orchestrator agent") + print(f"[ORCH_INIT] Orchestrator agent initialized successfully: {orchestrator_agent}") return orchestrator_agent diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index d92cf4e..4821f80 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -28,7 +28,7 @@ print("[INIT] Loading agents...") print(f"[INIT] Planner agent loaded: {planner_agent}") print(f"[INIT] Coder agent loaded: {coder_agent}") -print(f"[INIT] Orchestrator agent loaded: {get_orchestrator_agent}") +print(f"[INIT] Orchestrator agent loaded: {orchestrator_agent}") class DateTimeEncoder(json.JSONEncoder): """Custom JSON encoder that can handle datetime objects""" @@ -97,7 +97,7 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: print("Orchestrator agent running") - agent = orchestrator_agent(self.model_preference) + agent = await orchestrator_agent(self.model_preference) orchestrator_response = await agent.run( user_prompt=task, deps=deps_for_orchestrator diff --git a/ta-browser/core/server/routes/web.py b/ta-browser/core/server/routes/web.py index 8dabe34..1be01af 100644 --- a/ta-browser/core/server/routes/web.py +++ b/ta-browser/core/server/routes/web.py @@ -154,12 +154,3 @@ async def stream_session( logger.error(f"Error in /stream: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) -@router.get("/api/v1/model/preference") -async def get_model_preference(): - """Get the current model preference""" - try: - # Get model preference from environment or use default - model_preference = os.getenv("MODEL_PREFERENCE", "Anthropic") - return {"model_preference": model_preference} - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) \ No newline at end of file From 469c5165b6f8dd0b6cf5ee3de5a8389250032aa6 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Fri, 18 Apr 2025 15:36:46 +0530 Subject: [PATCH 18/24] feat: fixed websocket error --- cortex_on/instructor.py | 12 +++++++----- cortex_on/main.py | 17 +++++++++++++---- frontend/src/components/home/ChatList.tsx | 22 +++++++++++++++++++++- frontend/src/components/home/Header.tsx | 2 +- 4 files changed, 42 insertions(+), 11 deletions(-) diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index 4821f80..940d11c 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -128,14 +128,16 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: finally: logfire.info("Orchestration process complete") - # Clear any sensitive data + + async def shutdown(self): """Clean shutdown of orchestrator""" try: - # Close websocket if open - if self.websocket: - await self.websocket.close() - + #Close websocket if open + # if self.websocket: + # await self.websocket.close() + # logfire.info("websocket shutdown complete") + # print("websocket shutdown complete") # Clear all responses self.orchestrator_response = [] diff --git a/cortex_on/main.py b/cortex_on/main.py index 4488c3a..55926b1 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -95,7 +95,16 @@ async def websocket_endpoint(websocket: WebSocket): await websocket.accept() model_preference = app.state.model_preference print(f"[WEBSOCKET] New connection using model preference: {model_preference}") - while True: - data = await websocket.receive_text() - print(f"[WEBSOCKET] Received message, using model: {model_preference}") - await generate_response(data, websocket, model_preference) + try: + while True: + try: + data = await websocket.receive_text() + print(f"[WEBSOCKET] Received message, using model: {model_preference}") + await generate_response(data, websocket, model_preference) + except Exception as e: + print(f"[WEBSOCKET] Error handling message: {str(e)}") + break + except Exception as e: + print(f"[WEBSOCKET] Connection error: {str(e)}") + finally: + print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/ChatList.tsx b/frontend/src/components/home/ChatList.tsx index 9d8cb21..5c6bc7c 100644 --- a/frontend/src/components/home/ChatList.tsx +++ b/frontend/src/components/home/ChatList.tsx @@ -94,9 +94,29 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { retryOnError: true, shouldReconnect: () => true, reconnectInterval: 3000, + share: true, } ); + useEffect(() => { + if (readyState === ReadyState.CONNECTING) { + if (messages.length > 0) { + setIsLoading(true); + } + } else if (readyState === ReadyState.OPEN) { + setIsLoading(false); + } + }, [readyState, messages.length]); + + useEffect(() => { + if (messages.length === 0) { + setLiveUrl(""); + setCurrentOutput(null); + setOutputsList([]); + setIsLoading(false); + } + }, [messages.length]); + const scrollToBottom = (smooth = true) => { if (scrollAreaRef.current) { const scrollableDiv = scrollAreaRef.current.querySelector( @@ -557,7 +577,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { className="space-y-2 animate-fade-in animate-once animate-duration-500" key={idx} > - {readyState === ReadyState.CONNECTING ? ( + {readyState === ReadyState.CONNECTING && messages.length > 0 ? ( <> diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 5fba65a..d91c194 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - nav("/"); + window.location.href = "/"; }} > From be6bd57c43ae80d71c588f8b93bb32c4baa5fb8e Mon Sep 17 00:00:00 2001 From: aryan Date: Fri, 18 Apr 2025 17:12:39 +0530 Subject: [PATCH 19/24] fix(logging): Replace print statements with logfire for improved logging consistency across the application --- cortex_on/agents/orchestrator_agent.py | 2 +- cortex_on/agents/web_surfer.py | 2 +- cortex_on/instructor.py | 2 +- cortex_on/main.py | 31 ++++++++++++++----------- frontend/src/components/home/Header.tsx | 2 +- 5 files changed, 22 insertions(+), 17 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 2bfda6a..bdb177a 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -181,7 +181,7 @@ class orchestrator_deps: # deps_type=orchestrator_deps # ) # print("[ORCH_INIT] Orchestrator agent created successfully") -async def orchestrator_agent(model_preference: str) : +async def orchestrator_agent(model_preference: str) -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() elif model_preference == "OpenAI": diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 03503d5..2711348 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -47,7 +47,7 @@ async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, An "critique_disabled": False, "model_preference": self.model_preference } - print(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") + logfire.info(f"Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index 940d11c..388a7bb 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -41,7 +41,7 @@ def default(self, obj): # Main Orchestrator Class class SystemInstructor: def __init__(self, model_preference: str = "Anthropic"): - print(f"[INIT] Initializing SystemInstructor with model_preference: {model_preference}") + logfire.info(f"Initializing SystemInstructor with model_preference: {model_preference}") self.websocket: Optional[WebSocket] = None self.stream_output: Optional[StreamResponse] = None self.orchestrator_response: List[StreamResponse] = [] diff --git a/cortex_on/main.py b/cortex_on/main.py index 55926b1..deb1f38 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -3,8 +3,9 @@ from contextlib import asynccontextmanager # Third-party imports -from fastapi import FastAPI, WebSocket, Depends +from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends from fastapi.middleware.cors import CORSMiddleware +import logfire # Local application imports from instructor import SystemInstructor @@ -20,7 +21,7 @@ async def lifespan(app: FastAPI): # Set default model preference at startup app.state.model_preference = MODEL_PREFERENCE - print(f"[STARTUP] Setting default model preference to: {MODEL_PREFERENCE}") + logfire.info(f"Setting default model preference to: {MODEL_PREFERENCE}") # Initialize the instructor global instructor @@ -49,7 +50,7 @@ async def get_model_preference() -> str: """ Get the current model preference from app state """ - print(f"[GET] Current model preference: {app.state.model_preference}") + logfire.info(f"Current model preference: {app.state.model_preference}") return app.state.model_preference @app.get("/set_model_preference") @@ -58,10 +59,10 @@ async def set_model_preference(model: str): Set the model preference (Anthropic or OpenAI) and reinitialize the instructor """ if model not in ["Anthropic", "OpenAI"]: - print(f"[ERROR] Invalid model preference attempted: {model}") + logfire.error(f"Invalid model preference attempted: {model}") return {"error": "Model must be 'Anthropic' or 'OpenAI'"} - print(f"[SET] Changing model preference from {app.state.model_preference} to {model}") + logfire.info(f"Changing model preference from {app.state.model_preference} to {model}") app.state.model_preference = model # Reinitialize the instructor with new model preference @@ -69,14 +70,14 @@ async def set_model_preference(model: str): if instructor: await instructor.shutdown() instructor = SystemInstructor(model_preference=model) - print(f"[SET] Instructor reinitialized with model preference: {model}") + logfire.info(f"Instructor reinitialized with model preference: {model}") return {"message": f"Model preference set to {model}"} async def generate_response(task: str, websocket: Optional[WebSocket] = None, model_preference: str = None): if model_preference is None: model_preference = app.state.model_preference - print(f"[GENERATE] Using model preference: {model_preference} for task: {task[:30]}...") + logfire.info(f"Using model preference: {model_preference} for task: {task[:30]}...") global instructor if not instructor: @@ -86,7 +87,7 @@ async def generate_response(task: str, websocket: Optional[WebSocket] = None, mo @app.get("/agent/chat") async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List: - print(f"[CHAT] Received chat request with model preference: {model_preference}") + logfire.info(f"Received chat request with model preference: {model_preference}") final_agent_response = await generate_response(task, model_preference=model_preference) return final_agent_response @@ -94,17 +95,21 @@ async def agent_chat(task: str, model_preference: str = Depends(get_model_prefer async def websocket_endpoint(websocket: WebSocket): await websocket.accept() model_preference = app.state.model_preference - print(f"[WEBSOCKET] New connection using model preference: {model_preference}") + logfire.info(f"New connection using model preference: {model_preference}") try: while True: try: + # Use a timeout to detect if the client is still connected data = await websocket.receive_text() - print(f"[WEBSOCKET] Received message, using model: {model_preference}") + logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) + except WebSocketDisconnect: + print("[WEBSOCKET] Client disconnected") + break except Exception as e: print(f"[WEBSOCKET] Error handling message: {str(e)}") - break + if "disconnect message has been received" in str(e): + print("[WEBSOCKET] Disconnect detected, closing connection") + break except Exception as e: print(f"[WEBSOCKET] Connection error: {str(e)}") - finally: - print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index d91c194..5fba65a 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - window.location.href = "/"; + nav("/"); }} > From e3579cc2167c42f5b7c7191f100550813156081e Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Tue, 22 Apr 2025 22:35:44 +0530 Subject: [PATCH 20/24] feat[openai compatibility]:made changes that CortexON can now switches between OpenAI and Antropic models --- cortex_on/agents/orchestrator_agent.py | 49 +++++++--------------- cortex_on/agents/planner_agent.py | 6 +-- cortex_on/agents/web_surfer.py | 2 +- cortex_on/main.py | 13 +++--- frontend/src/components/home/ChatList.tsx | 3 +- frontend/src/components/home/Header.tsx | 2 +- ta-browser/core/utils/openai_msg_parser.py | 1 + 7 files changed, 27 insertions(+), 49 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index bdb177a..c344213 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -16,8 +16,6 @@ from agents.code_agent import coder_agent, CoderAgentDeps from utils.ant_client import get_client, get_anthropic_model_instance, get_openai_model_instance, get_openai_client -print("[ORCH_INIT] Starting orchestrator agent initialization") - @dataclass class orchestrator_deps: websocket: Optional[WebSocket] = None @@ -26,7 +24,6 @@ class orchestrator_deps: agent_responses: Optional[List[StreamResponse]] = None model_preference: str = "Anthropic" - orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. [AGENT CAPABILITIES] @@ -162,25 +159,6 @@ class orchestrator_deps: - Format: "Task description (agent_name)" """ -# print("[ORCH_INIT] Getting Anthropic client...") -# client = get_client() -# print(f"[ORCH_INIT] Anthropic client obtained: {client}") - -# print("[ORCH_INIT] Initializing Anthropic model...") -# model = get_anthropic_model_instance() - -# print(f"[ORCH_INIT] Anthropic model initialized: {model}") - -# print("Orchestrator agent model initialized without MODEL PREFERENCE") - -# print("[ORCH_INIT] Creating orchestrator agent...") -# orchestrator_agent = Agent( -# model=model, -# name="Orchestrator Agent", -# system_prompt=orchestrator_system_prompt, -# deps_type=orchestrator_deps -# ) -# print("[ORCH_INIT] Orchestrator agent created successfully") async def orchestrator_agent(model_preference: str) -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() @@ -195,7 +173,7 @@ async def orchestrator_agent(model_preference: str) -> Agent: system_prompt=orchestrator_system_prompt, deps_type=orchestrator_deps ) - + @orchestrator_agent.tool async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: """Plans the task and assigns it to the appropriate agents""" @@ -226,7 +204,7 @@ async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: planner_response = await agent.run(user_prompt=task) logfire.info(f"Planner Agent using model type: {ctx.deps.model_preference}") - + # Update planner stream with results plan_text = planner_response.data.plan planner_stream_output.steps.append("Task planned successfully") @@ -281,8 +259,7 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: # Create deps with the new stream_output deps_for_coder_agent = CoderAgentDeps( websocket=ctx.deps.websocket, - stream_output=coder_stream_output, - model_preference = ctx.deps.model_preference + stream_output=coder_stream_output ) # Run coder agent @@ -294,6 +271,7 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: logfire.info(f"coder_response: {coder_response}") logfire.info(f"Coder Agent using model type: {ctx.deps.model_preference}") + # Extract response data response_data = coder_response.data.content @@ -346,7 +324,7 @@ async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: model_preference=ctx.deps.model_preference ) logfire.info(f"web_surfing on model type: {ctx.deps.model_preference}") - + # Run WebSurfer with its own stream_output success, message, messages = await web_surfer_agent.generate_reply( instruction=task, @@ -479,9 +457,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas agent = await planner_agent(model_preference=ctx.deps.model_preference) plan_response = await agent.run(user_prompt=plan_prompt) current_content = plan_response.data.plan - logfire.info(f"Updating plan task using model type: {ctx.deps.model_preference}") - + else: # Read existing todo.md with open(todo_path, "r") as file: @@ -502,12 +479,12 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas planner_stream_output.steps.append("Asking planner to update the plan...") await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - agent = await planner_agent(model_preference=ctx.deps.model_preference) - updated_plan_response = await agent.run(user_prompt=update_prompt) + updated_plan_response = await planner_agent.run(user_prompt=update_prompt) updated_plan = updated_plan_response.data.plan + logfire.info(f"Updating prompt for plan task using model type: {ctx.deps.model_preference}") - + # Write the updated plan back to todo.md with open(todo_path, "w") as file: file.write(updated_plan) @@ -529,7 +506,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.error(error_msg, exc_info=True) planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = 500 + planner_stream_output.status_code = a500 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) return f"Failed to update the plan: {error_msg}" @@ -544,12 +521,16 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) return f"Failed to update plan: {error_msg}" - + + + logfire.info("All tools initialized for orchestrator agent") print(f"[ORCH_INIT] Orchestrator agent initialized successfully: {orchestrator_agent}") + return orchestrator_agent + # Helper function for sending WebSocket messages async def _safe_websocket_send(websocket: Optional[WebSocket], message: Any) -> bool: """Safely send message through websocket with error handling""" diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 5d29b7d..d5d7832 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -176,9 +176,6 @@ class PlannerResult(BaseModel): plan: str = Field(description="The generated or updated plan in string format - this should be the complete plan text") - - - async def planner_agent(model_preference: str = "Anthropic") -> Agent: if model_preference == "Anthropic": model = get_anthropic_model_instance() @@ -351,6 +348,7 @@ async def execute_terminal(command: str) -> str: except Exception as e: logfire.error(f"Error executing command: {str(e)}", exc_info=True) return f"Error executing command: {str(e)}" - + + logfire.info("All tools initialized for planner agent") return planner_agent \ No newline at end of file diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 2711348..e9fe772 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -47,7 +47,7 @@ async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, An "critique_disabled": False, "model_preference": self.model_preference } - logfire.info(f"Making API call with model_preference: {self.model_preference}") + logfire.info(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/main.py b/cortex_on/main.py index deb1f38..bb3aa93 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -3,14 +3,13 @@ from contextlib import asynccontextmanager # Third-party imports -from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends +from fastapi import FastAPI, WebSocket, WebSocketDisconnect ,Depends from fastapi.middleware.cors import CORSMiddleware import logfire # Local application imports from instructor import SystemInstructor - # Default model preference is Anthropic MODEL_PREFERENCE = "Anthropic" @@ -22,7 +21,7 @@ async def lifespan(app: FastAPI): # Set default model preference at startup app.state.model_preference = MODEL_PREFERENCE logfire.info(f"Setting default model preference to: {MODEL_PREFERENCE}") - + # Initialize the instructor global instructor instructor = SystemInstructor(model_preference=MODEL_PREFERENCE) @@ -99,17 +98,15 @@ async def websocket_endpoint(websocket: WebSocket): try: while True: try: - # Use a timeout to detect if the client is still connected data = await websocket.receive_text() logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) except WebSocketDisconnect: - print("[WEBSOCKET] Client disconnected") + print("Websocket disconnected") break except Exception as e: print(f"[WEBSOCKET] Error handling message: {str(e)}") - if "disconnect message has been received" in str(e): - print("[WEBSOCKET] Disconnect detected, closing connection") - break except Exception as e: print(f"[WEBSOCKET] Connection error: {str(e)}") + # finally: + # print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/ChatList.tsx b/frontend/src/components/home/ChatList.tsx index 5c6bc7c..50fba0a 100644 --- a/frontend/src/components/home/ChatList.tsx +++ b/frontend/src/components/home/ChatList.tsx @@ -117,6 +117,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { } }, [messages.length]); + const scrollToBottom = (smooth = true) => { if (scrollAreaRef.current) { const scrollableDiv = scrollAreaRef.current.querySelector( @@ -577,7 +578,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { className="space-y-2 animate-fade-in animate-once animate-duration-500" key={idx} > - {readyState === ReadyState.CONNECTING && messages.length > 0 ? ( + {readyState === ReadyState.CONNECTING && messages.length > 0 ? ( <> diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 5fba65a..d91c194 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - nav("/"); + window.location.href = "/"; }} > diff --git a/ta-browser/core/utils/openai_msg_parser.py b/ta-browser/core/utils/openai_msg_parser.py index 1458125..af57cdc 100644 --- a/ta-browser/core/utils/openai_msg_parser.py +++ b/ta-browser/core/utils/openai_msg_parser.py @@ -646,6 +646,7 @@ def get_conversation_history(self) -> List[ChatCompletionMessageParam]: return self.conversation_history +# In openai_msg_parser.py (or wherever ConversationStorage is defined) from datetime import datetime import os import json From 13f7358617c2216c04203216458de6d8465353d4 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Tue, 22 Apr 2025 22:50:04 +0530 Subject: [PATCH 21/24] feat[openai compatibility]:made changes that CortexON can now switches between OpenAI and Antropic models --- cortex_on/agents/orchestrator_agent.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index c344213..ec2188e 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -160,7 +160,8 @@ class orchestrator_deps: """ async def orchestrator_agent(model_preference: str) -> Agent: - if model_preference == "Anthropic": + if model_preference == "Anthropic": + model = get_anthropic_model_instance() elif model_preference == "OpenAI": model = get_openai_model_instance() @@ -259,7 +260,8 @@ async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: # Create deps with the new stream_output deps_for_coder_agent = CoderAgentDeps( websocket=ctx.deps.websocket, - stream_output=coder_stream_output + stream_output=coder_stream_output, + model_preference=ctx.deps.model_preference ) # Run coder agent @@ -479,7 +481,8 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas planner_stream_output.steps.append("Asking planner to update the plan...") await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - updated_plan_response = await planner_agent.run(user_prompt=update_prompt) + agent = await planner_agent(model_preference=ctx.deps.model_preference) + updated_plan_response = await agent.run(user_prompt=update_prompt) updated_plan = updated_plan_response.data.plan @@ -506,7 +509,7 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.error(error_msg, exc_info=True) planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = a500 + planner_stream_output.status_code = 500 await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) return f"Failed to update the plan: {error_msg}" From 074d8bb9e12e54b348923e295016d3d901b1dfe0 Mon Sep 17 00:00:00 2001 From: Soumyajit22theagentic Date: Tue, 22 Apr 2025 23:31:41 +0530 Subject: [PATCH 22/24] feat[openai compatibility]:made changes that CortexON can now switches between OpenAI and Antropic models --- cortex_on/main.py | 12 +++++++++--- frontend/src/components/home/Header.tsx | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cortex_on/main.py b/cortex_on/main.py index bb3aa93..6d06755 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -7,6 +7,9 @@ from fastapi.middleware.cors import CORSMiddleware import logfire +# Configure Logfire +logfire.configure() + # Local application imports from instructor import SystemInstructor @@ -102,11 +105,14 @@ async def websocket_endpoint(websocket: WebSocket): logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) except WebSocketDisconnect: - print("Websocket disconnected") + print("[WEBSOCKET] Client disconnected") break except Exception as e: print(f"[WEBSOCKET] Error handling message: {str(e)}") + if "disconnect message has been received" in str(e): + print(f"[WEBSOCKET] DIsconnect detected, closing connection: {str(e)}") + break except Exception as e: print(f"[WEBSOCKET] Connection error: {str(e)}") - # finally: - # print("[WEBSOCKET] connection closed") + # finally: + # print("[WEBSOCKET] connection closed") diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index d91c194..5fba65a 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -44,7 +44,7 @@ const Header = () => { className="rounded-xl" onClick={() => { dispatch(setMessages([])); - window.location.href = "/"; + nav("/"); }} > From 077201a0d31d2e5986940228813b9476ea825dfe Mon Sep 17 00:00:00 2001 From: aryan Date: Wed, 23 Apr 2025 17:32:17 +0530 Subject: [PATCH 23/24] refactor(Header): update header layout and styling for improved responsiveness and usability - Adjusted logo container size and alignment - Enhanced button and model toggle button styling for better visibility - Improved overall layout structure for better responsiveness --- frontend/src/components/home/Header.tsx | 54 +++++++++++++------------ 1 file changed, 29 insertions(+), 25 deletions(-) diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 5fba65a..3c212fc 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -13,20 +13,20 @@ const Header = () => { return (
-
{ - dispatch(setMessages([])); - nav("/"); - }} - > - Logo -
-
+
{ + dispatch(setMessages([])); + nav("/"); + }} + > + Logo +
+
nav("/vault")} - className={`h-full flex justify-center items-center cursor-pointer border-b-2 px-4 hover:border-[#BD24CA] ${ + className={`h-full flex justify-center items-center cursor-pointer border-b-2 -mb-[2px] px-4 hover:border-[#BD24CA] ${ location.includes("/vault") ? "border-[#BD24CA]" : "border-background" @@ -34,22 +34,26 @@ const Header = () => { >

Vault

- - {/* Model Toggle Button */} -
- +
+ {/* Model Toggle Button - Smaller size */} +
+ +
+ + +
); }; From e6e789a6b6ba30e8466b408d66c9b236708116e1 Mon Sep 17 00:00:00 2001 From: aryan Date: Tue, 29 Apr 2025 21:33:51 +0530 Subject: [PATCH 24/24] refactor(logging): replace print statements with logfire for consistent logging across the application --- cortex_on/agents/orchestrator_agent.py | 3 --- cortex_on/agents/planner_agent.py | 1 - cortex_on/agents/web_surfer.py | 2 +- cortex_on/instructor.py | 14 -------------- cortex_on/main.py | 14 ++++---------- frontend/src/components/home/Header.tsx | 2 +- ta-browser/core/orchestrator.py | 8 +++----- ta-browser/core/utils/init_client.py | 7 ------- 8 files changed, 9 insertions(+), 42 deletions(-) diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index ec2188e..e08c881 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -167,7 +167,6 @@ async def orchestrator_agent(model_preference: str) -> Agent: model = get_openai_model_instance() else: raise ValueError(f"Unknown model_preference: {model_preference}") - print(f"[ORCH_INIT] Creating orchestrator agent with model: {model}") orchestrator_agent = Agent( model=model, name="Orchestrator Agent", @@ -528,8 +527,6 @@ async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_tas logfire.info("All tools initialized for orchestrator agent") - print(f"[ORCH_INIT] Orchestrator agent initialized successfully: {orchestrator_agent}") - return orchestrator_agent diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index d5d7832..d607d04 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -183,7 +183,6 @@ async def planner_agent(model_preference: str = "Anthropic") -> Agent: model = get_openai_model_instance() else: raise ValueError(f"Unknown model_preference: {model_preference}") - print(f"[PLANNER_INIT] Creating planner agent with model: {model}") planner_agent = Agent( model=model, name="Planner Agent", diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index e9fe772..2711348 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -47,7 +47,7 @@ async def _make_api_call(self, instruction: str) -> Tuple[int, List[Dict[str, An "critique_disabled": False, "model_preference": self.model_preference } - logfire.info(f"[WEB_SURFER] Making API call with model_preference: {self.model_preference}") + logfire.info(f"Making API call with model_preference: {self.model_preference}") async with session.post(self.api_url, json=payload) as response: if response.status != 200: error_text = await response.text() diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index 388a7bb..bf9670f 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -25,11 +25,6 @@ load_dotenv() -print("[INIT] Loading agents...") -print(f"[INIT] Planner agent loaded: {planner_agent}") -print(f"[INIT] Coder agent loaded: {coder_agent}") -print(f"[INIT] Orchestrator agent loaded: {orchestrator_agent}") - class DateTimeEncoder(json.JSONEncoder): """Custom JSON encoder that can handle datetime objects""" def default(self, obj): @@ -47,7 +42,6 @@ def __init__(self, model_preference: str = "Anthropic"): self.orchestrator_response: List[StreamResponse] = [] self.model_preference = model_preference self._setup_logging() - print("[INIT] SystemInstructor initialization complete") def _setup_logging(self) -> None: """Configure logging with proper formatting""" @@ -94,8 +88,6 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: await self._safe_websocket_send(stream_output) stream_output.steps.append("Agents initialized successfully") await self._safe_websocket_send(stream_output) - - print("Orchestrator agent running") agent = await orchestrator_agent(self.model_preference) orchestrator_response = await agent.run( @@ -133,12 +125,6 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: async def shutdown(self): """Clean shutdown of orchestrator""" try: - #Close websocket if open - # if self.websocket: - # await self.websocket.close() - # logfire.info("websocket shutdown complete") - # print("websocket shutdown complete") - # Clear all responses self.orchestrator_response = [] logfire.info("Orchestrator shutdown complete") diff --git a/cortex_on/main.py b/cortex_on/main.py index 6d06755..44b840b 100644 --- a/cortex_on/main.py +++ b/cortex_on/main.py @@ -28,13 +28,9 @@ async def lifespan(app: FastAPI): # Initialize the instructor global instructor instructor = SystemInstructor(model_preference=MODEL_PREFERENCE) - print("[STARTUP] Instructor initialized") yield - - # # Cleanup - # if instructor: - # await instructor.shutdown() + app: FastAPI = FastAPI(lifespan=lifespan) @@ -105,14 +101,12 @@ async def websocket_endpoint(websocket: WebSocket): logfire.info(f"Received message, using model: {model_preference}") await generate_response(data, websocket, model_preference) except WebSocketDisconnect: - print("[WEBSOCKET] Client disconnected") + logfire.info("[WEBSOCKET] Client disconnected") break except Exception as e: - print(f"[WEBSOCKET] Error handling message: {str(e)}") + logfire.error(f"[WEBSOCKET] Error handling message: {str(e)}") if "disconnect message has been received" in str(e): print(f"[WEBSOCKET] DIsconnect detected, closing connection: {str(e)}") break except Exception as e: - print(f"[WEBSOCKET] Connection error: {str(e)}") - # finally: - # print("[WEBSOCKET] connection closed") + logfire.error(f"[WEBSOCKET] Connection error: {str(e)}") diff --git a/frontend/src/components/home/Header.tsx b/frontend/src/components/home/Header.tsx index 3c212fc..4d142fb 100644 --- a/frontend/src/components/home/Header.tsx +++ b/frontend/src/components/home/Header.tsx @@ -26,7 +26,7 @@ const Header = () => {
nav("/vault")} - className={`h-full flex justify-center items-center cursor-pointer border-b-2 -mb-[2px] px-4 hover:border-[#BD24CA] ${ + className={`h-full flex justify-center items-center cursor-pointer border-b-2 -mb-[27px] px-4 pb-7 hover:border-[#BD24CA] ${ location.includes("/vault") ? "border-[#BD24CA]" : "border-background" diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 69df1d7..4d56511 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -252,7 +252,7 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com", m "start_url": start_url, "model_preference": model_preference }) - print(f"[ORCHESTRATOR] Starting async_init with model_preference: {model_preference}") + logger.info(f"Starting async_init with model_preference: {model_preference}") # Store and validate parameters self.job_id = job_id @@ -283,15 +283,13 @@ async def async_init(self, job_id: str, start_url: str = "https://google.com", m # 4. Initialize client and agents try: - print(f"[ORCHESTRATOR] Calling initialize_client with model_preference: {model_preference}") from core.utils.init_client import initialize_client self.client, model_instance = await initialize_client(model_preference) - print(f"[ORCHESTRATOR] Client initialized successfully, now initializing agents") self.initialize_agents(model_instance) - print(f"[ORCHESTRATOR] Agents initialized successfully with model_preference: {model_preference}") + logger.info(f"Agents initialized successfully with model_preference: {model_preference}") except Exception as agent_error: error_msg = f"Failed to initialize client and agents: {str(agent_error)}" - print(f"[ORCHESTRATOR] ERROR: {error_msg}") + logger.error(f"{error_msg}") raise RuntimeError(error_msg) from agent_error self.async_init_done = True diff --git a/ta-browser/core/utils/init_client.py b/ta-browser/core/utils/init_client.py index 5a423b9..5be2ea3 100644 --- a/ta-browser/core/utils/init_client.py +++ b/ta-browser/core/utils/init_client.py @@ -16,7 +16,6 @@ async def initialize_client(model_preference: str = "Anthropic"): tuple: (client_instance, model_instance) """ try: - print(f"[INIT_CLIENT] *** Initializing client with model preference: {model_preference} ***") logger.info(f"Initializing client with model preference: {model_preference}") if model_preference == "Anthropic": # Get API key from environment variable @@ -27,7 +26,6 @@ async def initialize_client(model_preference: str = "Anthropic"): # Set model name - Claude 3.5 Sonnet model_name = os.getenv("ANTHROPIC_MODEL_NAME") - print(f"[INIT_CLIENT] Using Anthropic model: {model_name}") # Create client config config = { @@ -44,7 +42,6 @@ async def initialize_client(model_preference: str = "Anthropic"): from pydantic_ai.models.anthropic import AnthropicModel model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) - print(f"[INIT_CLIENT] Anthropic client initialized successfully with model: {model_name}") logger.info(f"Anthropic client initialized successfully with model: {model_name}") return client_instance, model_instance elif model_preference == "OpenAI": @@ -56,7 +53,6 @@ async def initialize_client(model_preference: str = "Anthropic"): # Set model name - GPT-4o model_name = os.getenv("OPENAI_MODEL_NAME") - print(f"[INIT_CLIENT] Using OpenAI model: {model_name}") # Create client config config = { @@ -73,16 +69,13 @@ async def initialize_client(model_preference: str = "Anthropic"): from pydantic_ai.models.openai import OpenAIModel model_instance = OpenAIModel(model_name=model_name, openai_client=client_instance) - print(f"[INIT_CLIENT] OpenAI client initialized successfully with model: {model_name}") logger.info(f"OpenAI client initialized successfully with model: {model_name}") return client_instance, model_instance else: error_msg = f"Invalid model preference: {model_preference}. Must be 'Anthropic' or 'OpenAI'" - print(f"[INIT_CLIENT] ERROR: {error_msg}") raise ValueError(error_msg) except Exception as e: error_msg = f"Error initializing client: {str(e)}" - print(f"[INIT_CLIENT] CRITICAL ERROR: {error_msg}") logger.error(error_msg, exc_info=True) raise \ No newline at end of file