Update app.py

app.py

@@ -1,49 +1,63 @@
 import torch
+import requests
 from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 from transformers import pipeline
 
 main = FastAPI()
 
-main.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
+# 🔱 Inachi Core Configuration
 MODEL_ID = "google/gemma-3-1b-it"
 
-# 🔱 Optimization 1: KV Cache භාවිතය (වේගය වැඩි කිරීමට)
+# Load the model with CPU optimizations
 pipe = pipeline(
     "text-generation",
     model=MODEL_ID,
     device_map="cpu",
-    torch_dtype=torch.bfloat16, # CPU එක bfloat16 වලට කැමතියි
-    use_cache=True, # කලින් ජෙනරේට් කරපු ටෝකන් මතක තබා ගනී
+    torch_dtype=torch.bfloat16,
     trust_remote_code=True
 )
 
-class ChatRequest(BaseModel):
-    message: str
+# 🔱 Simple Web Search Tool (DuckDuckGo API - No Key Required)
+def web_search(query):
+    try:
+        url = f"https://api.duckduckgo.com/?q={query}&format=json"
+        response = requests.get(url).json()
+        return response.get("AbstractText", "No specific web data found.")
+    except:
+        return "Search failed."
 
 @main.post("/v1/chat")
-async def chat(request_data: ChatRequest):
-    user_query = request_data.message.strip()
+async def chat(data: dict):
+    user_query = data.get("message", "")
+    
+    # 🔱 Inachi Identity & System Instruction
+    # මෙතනින් තමයි එයාට තමන් කවුද කියලා කියලා දෙන්නේ
+    system_prompt = (
+        "You are Inachi AI, a highly advanced assistant developed by the Inachi Team. "
+        "Your goal is to provide technical, precise, and helpful information. "
+        "Always identify yourself as Inachi AI when asked."
+    )
     
-    messages = [
-        {"role": "user", "content": user_query},
-    ]
+    # Web search check (සරලව search කරන්න ඕනෙද කියලා බලනවා)
+    search_context = ""
+    if "search" in user_query.lower() or "latest" in user_query.lower():
+        search_context = f"\nWeb Search Result: {web_search(user_query)}"
+
+    # Prompt එක සැකසීම
+    full_prompt = f"{system_prompt}\nContext: {search_context}\nUser: {user_query}\nInachi AI:"
     
-    # 🔱 Optimization 2: Coding සඳහා ප්‍රමාණවත් ඉඩක් ලබා දීම
     results = pipe(
-        messages, 
-        max_new_tokens=1024, # 🔱 දැන් ඔයාට දිග කෝඩ් එකක් වුණත් ගන්න පුළුවන්
-        do_sample=False,      # Coding වලට sample ඕනේ නැහැ, Greedy search එක වේගවත්
-        temperature=0.0,     # වඩාත් නිවැරදි කෝඩ් එකක් සඳහා (Zero randomness)
-        pad_token_id=50256
+        full_prompt, 
+        max_new_tokens=512, 
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
     )
+    
+    # පිරිසිදු පිළිතුර ලබා ගැනීම
+    reply = results[0]['generated_text'].split("Inachi AI:")[-1].strip()
+    return {"reply": reply}
 
-    reply = results[0]['generated_text'][-1]['content']
-    return {"reply": reply}
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(main, host="0.0.0.0", port=7860)