Feature/model preference

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
+                }
+                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()

cortex_on/instructor.py

@@ -25,9 +25,6 @@
 
 load_dotenv()
 
-
-
-
 class DateTimeEncoder(json.JSONEncoder):
     """Custom JSON encoder that can handle datetime objects"""
     def default(self, obj):
@@ -38,10 +35,12 @@ def default(self, obj):
 
 # Main Orchestrator Class
 class SystemInstructor:
-    def __init__(self):
+    def __init__(self, model_preference: str = "Anthropic"):
+        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] = []
+        self.model_preference = model_preference
         self._setup_logging()
 
     def _setup_logging(self) -> None:
@@ -80,19 +79,24 @@ 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:
             # Initialize system
             await self._safe_websocket_send(stream_output)
             stream_output.steps.append("Agents initialized successfully")
             await self._safe_websocket_send(stream_output)
-
-            orchestrator_response = await orchestrator_agent.run(
+
+            agent = await 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}")
@@ -116,15 +120,11 @@ 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()
-
-            # Clear all responses
             self.orchestrator_response = []
 
             logfire.info("Orchestrator shutdown complete")

cortex_on/main.py

@@ -1,27 +1,112 @@
 # 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, WebSocketDisconnect ,Depends
+from fastapi.middleware.cors import CORSMiddleware
+import logfire
+
+# Configure Logfire
+logfire.configure()
 
 # Local application imports
 from instructor import SystemInstructor
 
+# 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
+    logfire.info(f"Setting default model preference to: {MODEL_PREFERENCE}")
+
+    # Initialize the instructor
+    global instructor
+    instructor = SystemInstructor(model_preference=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
+    """
+    logfire.info(f"Current model preference: {app.state.model_preference}")
+    return app.state.model_preference
 
-app: FastAPI = FastAPI()
+@app.get("/set_model_preference")
+async def set_model_preference(model: str):
+    """
+    Set the model preference (Anthropic or OpenAI) and reinitialize the instructor
+    """
+    if model not in ["Anthropic", "OpenAI"]:
+        logfire.error(f"Invalid model preference attempted: {model}")
+        return {"error": "Model must be 'Anthropic' or 'OpenAI'"}
+
+    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
+    global instructor
+    if instructor:
+        await instructor.shutdown()
+    instructor = SystemInstructor(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):
-    orchestrator: SystemInstructor = SystemInstructor()
-    return await orchestrator.run(task, websocket)
+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
+    logfire.info(f"Using model preference: {model_preference} for task: {task[:30]}...")
+
+    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) -> List:
-    final_agent_response = await generate_response(task)
+async def agent_chat(task: str, model_preference: str = Depends(get_model_preference)) -> List:
+    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
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
-    while True:
-        data = await websocket.receive_text()
-        await generate_response(data, websocket)
+    model_preference = app.state.model_preference
+    logfire.info(f"New connection using model preference: {model_preference}")
+    try:
+        while True:
+            try:
+                data = await websocket.receive_text()
+                logfire.info(f"Received message, using model: {model_preference}")
+                await generate_response(data, websocket, model_preference)
+            except WebSocketDisconnect:
+                logfire.info("[WEBSOCKET] Client disconnected")
+                break
+            except Exception as 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:
+        logfire.error(f"[WEBSOCKET] Connection error: {str(e)}")

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

frontend/src/components/home/ChatList.tsx

@@ -94,9 +94,30 @@ 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 +578,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  ? (
                     <>
                       <Skeleton className="h-[3vh] w-[10%] animate-pulse" />
                       <Skeleton className="h-[2vh] w-[80%] animate-pulse animate-delay-100" />

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();
@@ -12,19 +13,20 @@ const Header = () => {
 
   return (
     <div className="h-[8vh] border-b-2 flex justify-between items-center px-8">
-      <div
-        className="w-[12%]"
-        onClick={() => {
-          dispatch(setMessages([]));
-          nav("/");
-        }}
-      >
-        <img src={Logo} alt="Logo" />
-      </div>
-      <div className="w-full h-full gap-2 items-center px-4">
+      <div className="flex items-center gap-6">
+        <div
+          className="cursor-pointer w-[140px]"
+          onClick={() => {
+            dispatch(setMessages([]));
+            nav("/");
+          }}
+        >
+          <img src={Logo} alt="Logo" className="w-full" />
+        </div>
+
         <div
           onClick={() => 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 -mb-[27px] px-4 pb-7 hover:border-[#BD24CA] ${
             location.includes("/vault")
               ? "border-[#BD24CA]"
               : "border-background"
@@ -33,17 +35,25 @@ const Header = () => {
           <p className="text-xl font-medium">Vault</p>
         </div>
       </div>
-      <Button
-        size="sm"
-        className="rounded-xl"
-        onClick={() => {
-          dispatch(setMessages([]));
-          nav("/");
-        }}
-      >
-        <MessageCirclePlus size={20} absoluteStrokeWidth />
-        New Chat
-      </Button>
+
+      <div className="flex items-center gap-4">
+        {/* Model Toggle Button - Smaller size */}
+        <div className="scale-75 hover:opacity-90 transition-opacity cursor-pointer">
+          <ModelToggle />
+        </div>
+
+        <Button
+          size="sm"
+          className="rounded-xl"
+          onClick={() => {
+            dispatch(setMessages([]));
+            nav("/");
+          }}
+        >
+          <MessageCirclePlus size={20} absoluteStrokeWidth />
+          New Chat
+        </Button>
+      </div>
     </div>
   );
 };