Update main.py

main.py

@@ -12,17 +12,7 @@ async def startup():
 
 from pydantic import BaseModel
 from typing import List, Dict, Any
-from helper_functions_api import md_to_html, search_brave, fetch_and_extract_content, limit_tokens, together_response, insert_data
-import os
-from dotenv import load_dotenv, find_dotenv
-from datetime import datetime, timedelta
-import asyncio
-import re
-# Load environment variables from .env file
-#load_dotenv("keys.env")
-
 
-    
 TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
 BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
@@ -30,95 +20,103 @@ HELICON_API_KEY = os.getenv("HELICON_API_KEY")
 SUPABASE_USER = os.environ['SUPABASE_USER']
 SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']
 
-# Groq model names
-llm_default_small = "llama3-8b-8192"
-llm_default_medium = "llama3-70b-8192"
-
-# Together Model names (fallback)
-llm_fallback_small = "meta-llama/Llama-3-8b-chat-hf"
-llm_fallback_medium = "meta-llama/Llama-3-70b-chat-hf"
-
-SysPromptJson = "You are now in the role of an expert AI who can extract structured information from user request. Both key and value pairs must be in double quotes. You must respond ONLY with a valid JSON file. Do not add any additional comments."
-SysPromptList = "You are now in the role of an expert AI who can extract structured information from user request. All elements must be in double quotes. You must respond ONLY with a valid python List. Do not add any additional comments."
-SysPromptDefault = "You are an expert AI, complete the given task. Do not add any additional comments."
-SysPromptMd = "You are an expert AI who can create a structured report using information provided in the context from user request.The report should be in markdown format consists of markdown tables structured into subtopics. Do not add any additional comments."
-
-sys_prompts = {
-    "offline": {
-        "Chat": "You are an expert AI, complete the given task. Do not add any additional comments.",
-        "Full Text Report": "You are an expert AI who can create a detailed report from user request. The report should be in markdown format. Do not add any additional comments.",
-        "Tabular Report": "You are an expert AI who can create a structured report from user request.The report should be in markdown format structured into subtopics/tables/lists. Do not add any additional comments.",
-        "Tables only": "You are an expert AI who can create a structured tabular report from user request.The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
-    },
-    "online": {
-        "Chat": "You are an expert AI, complete the given task using the provided context. Do not add any additional comments.",
-        "Full Text Report": "You are an expert AI who can create a detailed report using information scraped from the internet. You should decide which information is relevant to the given task and use it to create a report. The report should be in markdown format. Do not add any additional comments.",
-        "Tabular Report": """You are an expert AI who can provide answers using internet search results.
-                                1 filter and summarize relevant information, if there are conflicting information, use the latest source.
-                                2. use it to construct a clear and factual answer.
-                                Your response should be properly formatted and well readable using markdown formatting. """,
-        "Tables only": "You are an expert AI who can create a structured tabular report using information scraped from the internet. You should decide which information is relevant to the given task. The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
-    },
-}
-
-class QueryModel(BaseModel):
-    topic: str = Query(default="", description="input query to generate Report")
-    description: str = Query(default="", description="additional context for report")
-    user_id: str = Query(default="", description="unique user id")
-    user_name: str = Query(default="", description="user name")
-    internet: bool = Query(default=True, description="Enable Internet search")
-    output_format: str = Query(default="Tabular Report", description="Output format for the report",
-                               enum=["Chat", "Full Text Report", "Tabular Report", "Tables only"])
-    data_format: str = Query(default="Structured data", description="Type of data to extract from the internet",
-                             enum=["No presets", "Structured data", "Quantitative data"])
-
-#@cache(expire=604800)
-async def generate_report(query: QueryModel):
-    query_str = query.topic
-    description = query.description
-    user_id = query.user_id
-    internet = "online" if query.internet else "offline"
-    sys_prompt_output_format = sys_prompts[internet][query.output_format]
-    data_format = query.data_format
-    optimized_search_query = ""
-    all_text_with_urls = [("", "")]
-
-    if query.internet:
-        search_query = re.sub(r'[^\w\s]', '', description).strip()
-        try:
-            urls, optimized_search_query = await search_brave(search_query, num_results=8)
-            all_text_with_urls = fetch_and_extract_content(data_format, urls, description)
-            additional_context = limit_tokens(str(all_text_with_urls))
-            prompt = f"Write a factual answer to the query:: {description} #### , refer the provided internet search results reference:{additional_context}"
-        except Exception as e:
-            print(e)
-            query.internet = False
-            print("failed to search/scrape results, falling back to LLM response")
-
-    if not query.internet:
-        prompt = f"#### COMPLETE THE TASK: {description} #### IN THE CONTEXT OF ### CONTEXT: {query_str}"
-
-    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=sys_prompt_output_format)
-
-    if user_id != "test":
-        insert_data(user_id, query_str, description, str(all_text_with_urls), md_report)
-    references_html = {}
-    for text, url in all_text_with_urls:
-        references_html[url] = str(md_to_html(text))
-
-    return {
-        "report": md_to_html(md_report),
-        "references": references_html,
-        "search_query": optimized_search_query
-    }
-
-@app.post("/generate_report")
-async def api_generate_report(request: Request, query: QueryModel):
-    return await generate_report(query)
-    
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
-    allow_headers=["*"],)
+    allow_headers=["*"],)
+
+import asyncio
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from jupyter_client import KernelManager
+from typing import Dict
+from datetime import datetime, timedelta
+import psutil
+
+app = FastAPI()
+
+# Model for input data
+class CodeExecutionRequest(BaseModel):
+    session_token: str
+    code: str
+
+# Store kernel managers and last access times
+kernel_managers: Dict[str, KernelManager] = {}
+last_access_times: Dict[str, datetime] = {}
+
+# Timeout duration in seconds
+TIMEOUT_DURATION = 600  # 10 minutes
+
+# Function to create a new kernel
+async def create_kernel(session_token: str):
+    km = KernelManager()
+    km.start_kernel()
+    kernel_managers[session_token] = km
+    last_access_times[session_token] = datetime.now()
+
+# Function to kill a kernel
+async def kill_kernel(session_token: str):
+    km = kernel_managers.pop(session_token, None)
+    if km:
+        km.shutdown_kernel(now=True)
+    last_access_times.pop(session_token, None)
+
+# Function to execute code in a kernel
+async def execute_code(session_token: str, code: str):
+    if session_token not in kernel_managers:
+        await create_kernel(session_token)
+
+    km = kernel_managers[session_token]
+    kc = km.client()
+    kc.start_channels()
+
+    try:
+        kc.execute(code)
+        reply = kc.get_shell_msg(timeout=5)
+        result = reply['content']
+    except Exception as e:
+        result = {'status': 'error', 'ename': type(e).__name__, 'evalue': str(e)}
+
+    kc.stop_channels()
+    last_access_times[session_token] = datetime.now()
+    return result
+
+# Background task to check for idle kernels
+async def check_idle_kernels():
+    while True:
+        now = datetime.now()
+        for session_token, last_access in list(last_access_times.items()):
+            if now - last_access > timedelta(seconds=TIMEOUT_DURATION):
+                await kill_kernel(session_token)
+        await asyncio.sleep(60)  # Check every minute
+
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(check_idle_kernels())
+
+@app.post("/execute")
+async def execute(request: CodeExecutionRequest):
+    result = await execute_code(request.session_token, request.code)
+    return result
+
+@app.get("/info")
+async def get_info():
+    # Get the number of active kernels
+    active_kernels = len(kernel_managers)
+
+    # Get system resource usage
+    cpu_usage = psutil.cpu_percent(interval=1)
+    ram_usage = psutil.virtual_memory().percent
+
+    # Return the information
+    return {
+        "active_kernels": active_kernels,
+        "cpu_usage_percent": cpu_usage,
+        "ram_usage_percent": ram_usage
+    }
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)