main.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319
  1. import json
  2. import logging
  3. import httpx
  4. import bcrypt
  5. from contextlib import asynccontextmanager
  6. from fastapi import FastAPI, HTTPException, Depends, Header
  7. from database import create_tables, create_user, get_user_by_username, create_session, get_user_from_token, delete_session, search_foods_by_name, save_chat_message, get_user_chat_history, get_user_profile, get_food_by_id
  8. from fastapi.responses import HTMLResponse, StreamingResponse
  9. from fastapi.staticfiles import StaticFiles
  10. from pydantic import BaseModel
  11. from typing import List, Generator, Optional
  12. logging.basicConfig(level=logging.INFO)
  13. logger = logging.getLogger(__name__)
  14. @asynccontextmanager
  15. async def lifespan(app: FastAPI):
  16. create_tables()
  17. yield
  18. app = FastAPI(title="LocalFoodAI Chat", lifespan=lifespan)
  19. # Use direct bcrypt for better environment compatibility
  20. def get_password_hash(password: str):
  21. # Hash requires bytes
  22. pwd_bytes = password.encode('utf-8')
  23. salt = bcrypt.gensalt()
  24. hashed = bcrypt.hashpw(pwd_bytes, salt)
  25. return hashed.decode('utf-8')
  26. def verify_password(plain_password: str, hashed_password: str):
  27. # bcrypt.checkpw handles verification
  28. return bcrypt.checkpw(plain_password.encode('utf-8'), hashed_password.encode('utf-8'))
  29. class UserCreate(BaseModel):
  30. username: str
  31. password: str
  32. class UserLogin(BaseModel):
  33. username: str
  34. password: str
  35. async def get_current_user(authorization: Optional[str] = Header(None)):
  36. if not authorization or not authorization.startswith("Bearer "):
  37. raise HTTPException(status_code=401, detail="Authentication required")
  38. token = authorization.split(" ")[1]
  39. user = get_user_from_token(token)
  40. if not user:
  41. raise HTTPException(status_code=401, detail="Invalid or expired session")
  42. return user
  43. OLLAMA_URL = "http://localhost:11434/api/chat"
  44. MODEL_NAME = "llama3.1:8b"
  45. # Common stopwords to strip before searching the food database
  46. _STOPWORDS = {
  47. 'how', 'many', 'much', 'calories', 'does', 'have', 'has', 'is', 'are',
  48. 'in', 'the', 'a', 'an', 'of', 'for', 'with', 'what', 'tell', 'me',
  49. 'about', 'nutritional', 'value', 'nutrition', 'macro', 'macros',
  50. 'protein', 'fat', 'carbs', 'fiber', 'can', 'you', 'i', 'want', 'need',
  51. 'eat', 'eating', 'food', 'meal', 'diet', 'healthy', 'make', 'cook',
  52. 'recipe', 'per', '100g', 'gram', 'grams', 'serving'
  53. }
  54. def extract_food_context(messages: list) -> str | None:
  55. """Scan the last user message for food keywords and enrich with local DB data."""
  56. # Find the last user message
  57. last_user_msg = None
  58. for msg in reversed(messages):
  59. role = msg.get('role', '') if isinstance(msg, dict) else msg.role
  60. content = msg.get('content', '') if isinstance(msg, dict) else msg.content
  61. if role == 'user':
  62. last_user_msg = content
  63. break
  64. if not last_user_msg:
  65. return None
  66. # Extract meaningful keywords by removing stopwords
  67. words = last_user_msg.lower().replace('?', '').replace(',', '').split()
  68. keywords = [w for w in words if w not in _STOPWORDS and len(w) > 2]
  69. if not keywords:
  70. return None
  71. # Try each keyword against the local food database, collect unique results
  72. found_items = {}
  73. # Optimization: Only use the first 2 most relevant keywords to keep context small on CPU
  74. for kw in keywords[:2]:
  75. results = search_foods_by_name(kw, limit=2)
  76. for item in results:
  77. # Truncate extremely long USDA names for performance
  78. short_name = item['name'][:100] + ("..." if len(item['name']) > 100 else "")
  79. if short_name not in found_items:
  80. found_items[short_name] = item
  81. if len(found_items) >= 3:
  82. break
  83. if not found_items:
  84. return None
  85. # Build a structured context block for the system prompt
  86. lines = [
  87. "[SYSTEM: NUTRITIONAL ANALYST MODE]",
  88. "You are the LocalFoodAI Analyst. Use ONLY the following verified local data for nutritional values.",
  89. "Always prioritize this data over your training memory. If a value is 0.0, it means it was not found or is negligible.",
  90. "If a specific nutrient is not listed below, state that it is not available in the local database.",
  91. "Format your summary clearly (tables are preferred).",
  92. ""
  93. ]
  94. for name, item in found_items.items():
  95. # Compact, token-efficient format for the LLM
  96. line = (
  97. f"- {name}: {item['calories']}kcal | P:{item['protein_g']}g | F:{item['fat_g']}g | C:{item['carbs_g']}g | "
  98. f"Fib:{item['fiber_g']}g | Sug:{item['sugar_g']}g | Chol:{item['cholesterol_mg']}mg | "
  99. f"Na:{item['sodium_mg']}mg | Ca:{item['calcium_mg']}mg | Fe:{item['iron_mg']}mg | "
  100. f"K:{item['potassium_mg']}mg | VitA:{item['vitamin_a_iu']}IU | VitC:{item['vitamin_c_mg']}mg"
  101. )
  102. lines.append(line)
  103. return "\n".join(lines)
  104. # Mount static files to serve the frontend
  105. app.mount("/static", StaticFiles(directory="static"), name="static")
  106. class ChatMessage(BaseModel):
  107. role: str
  108. content: str
  109. class ChatRequest(BaseModel):
  110. messages: List[ChatMessage]
  111. @app.get("/", response_class=HTMLResponse)
  112. async def read_root():
  113. """Serve the chat interface HTML"""
  114. try:
  115. with open("static/index.html", "r", encoding="utf-8") as f:
  116. return HTMLResponse(content=f.read())
  117. except FileNotFoundError:
  118. return HTMLResponse(content="<h1>Welcome to LocalFoodAI</h1><p>static/index.html not found. Please create the frontend.</p>")
  119. @app.post("/api/register")
  120. async def register_user(user: UserCreate):
  121. if len(user.username.strip()) < 3:
  122. raise HTTPException(status_code=400, detail="Username must be at least 3 characters")
  123. if len(user.password.strip()) < 6:
  124. raise HTTPException(status_code=400, detail="Password must be at least 6 characters")
  125. hashed_password = get_password_hash(user.password)
  126. user_id = create_user(user.username.strip(), hashed_password)
  127. if not user_id:
  128. raise HTTPException(status_code=400, detail="Username already exists")
  129. # Auto-login after registration
  130. token = create_session(user_id)
  131. return {"message": "User registered successfully", "token": token, "username": user.username.strip()}
  132. @app.post("/api/login")
  133. async def login_user(user: UserLogin):
  134. db_user = get_user_by_username(user.username.strip())
  135. if not db_user:
  136. raise HTTPException(status_code=401, detail="Invalid username or password")
  137. if not verify_password(user.password, db_user["password_hash"]):
  138. raise HTTPException(status_code=401, detail="Invalid username or password")
  139. token = create_session(db_user["id"])
  140. return {"status": "success", "username": db_user["username"], "token": token}
  141. @app.post("/api/logout")
  142. async def logout(authorization: Optional[str] = Header(None)):
  143. if authorization and authorization.startswith("Bearer "):
  144. token = authorization.split(" ")[1]
  145. delete_session(token)
  146. return {"message": "Logged out successfully"}
  147. @app.get("/api/macros/targets")
  148. async def get_macro_targets(current_user: dict = Depends(get_current_user)):
  149. """API endpoint to securely fetch the user's current macronutrient targets"""
  150. profile = get_user_profile(current_user['id'])
  151. if not profile:
  152. # Fallback to defaults in case database insertion failed
  153. return {
  154. "calories": 2000,
  155. "protein_g": 150,
  156. "carbs_g": 200,
  157. "fat_g": 65
  158. }
  159. return {
  160. "calories": profile.get("target_calories", 2000),
  161. "protein_g": profile.get("target_protein_g", 150),
  162. "carbs_g": profile.get("target_carbs_g", 200),
  163. "fat_g": profile.get("target_fat_g", 65)
  164. }
  165. @app.post("/chat")
  166. async def chat_endpoint(request: ChatRequest, current_user: dict = Depends(get_current_user)):
  167. """Proxy chat requests to the local Ollama instance with streaming support.
  168. Automatically enriches prompts with verified local SQLite nutritional data.
  169. """
  170. # Keep only the last 6 messages for context window performance on CPU
  171. all_messages = [msg.model_dump() for msg in request.messages]
  172. messages = all_messages[-6:]
  173. # Save the latest user message to DB
  174. if messages and messages[-1]['role'] == 'user':
  175. save_chat_message(current_user['id'], 'user', messages[-1]['content'])
  176. # --- TG-35: Local SQL RAG Enrichment ---
  177. db_context = extract_food_context(messages)
  178. if db_context:
  179. # Prepend as a system message so it acts as grounded knowledge
  180. # We ensure it's a short, concise instruction to prevent context bloat
  181. messages = [{"role": "system", "content": db_context}] + messages
  182. logger.info(f"[Chat] User '{current_user['username']}' is chatting. Context items: {'Yes' if db_context else 'No'}. Message count: {len(messages)}")
  183. payload = {
  184. "model": MODEL_NAME,
  185. "messages": messages,
  186. "stream": True
  187. }
  188. async def generate_response():
  189. try:
  190. bot_full_response = ""
  191. async with httpx.AsyncClient(timeout=300.0) as client:
  192. # Use a combined timeout for the entire request
  193. async with client.stream("POST", OLLAMA_URL, json=payload, timeout=300.0) as response:
  194. if response.status_code != 200:
  195. error_detail = await response.aread()
  196. logger.error(f"Ollama returned error {response.status_code}: {error_detail}")
  197. yield f"data: {json.dumps({'error': f'LLM Error ({response.status_code})'})}\n\n"
  198. return
  199. async for line in response.aiter_lines():
  200. if line:
  201. try:
  202. data = json.loads(line)
  203. if "message" in data and "content" in data["message"]:
  204. content = data["message"]["content"]
  205. bot_full_response += content
  206. yield f"data: {json.dumps({'content': content})}\n\n"
  207. if data.get("done"):
  208. break
  209. except json.JSONDecodeError:
  210. continue
  211. # Save final bot response to DB
  212. if bot_full_response.strip():
  213. save_chat_message(current_user['id'], 'assistant', bot_full_response)
  214. except Exception as e:
  215. logger.exception(f"Unexpected error during chat stream: {e}")
  216. yield f"data: {json.dumps({'error': 'A technical error occurred while generating the response.'})}\n\n"
  217. return StreamingResponse(generate_response(), media_type="text/event-stream")
  218. @app.get("/api/chat/history")
  219. async def get_history(current_user: dict = Depends(get_current_user)):
  220. """Fetch the chat history for the authenticated user"""
  221. history = get_user_chat_history(current_user['id'])
  222. return {"history": history}
  223. @app.get("/api/food/search")
  224. async def search_food(q: str, current_user: dict = Depends(get_current_user)):
  225. """API endpoint to search for food items securely using token authentication"""
  226. if not q or len(q.strip()) < 1:
  227. return {"results": []}
  228. logger.info(f"User {current_user['username']} searched for [{q}]")
  229. results = search_foods_by_name(q.strip(), limit=15)
  230. return {"results": results}
  231. @app.get("/api/food/{food_id}")
  232. async def get_food_detail(food_id: int, current_user: dict = Depends(get_current_user)):
  233. """API endpoint to fetch structured nutritional details for a specific food item"""
  234. food = get_food_by_id(food_id)
  235. if not food:
  236. raise HTTPException(status_code=404, detail="Food item not found")
  237. # Structure the data as proposed in the implementation plan
  238. structured_data = {
  239. "id": food["id"],
  240. "name": food["name"],
  241. "category": food["category"],
  242. "base_weight_g": food["base_weight_g"],
  243. "macros": {
  244. "calories": food["calories"],
  245. "protein_g": food["protein_g"],
  246. "fat_g": food["fat_g"],
  247. "carbs_g": food["carbs_g"]
  248. },
  249. "extended": {
  250. "fiber_g": food["fiber_g"],
  251. "sugar_g": food["sugar_g"],
  252. "cholesterol_mg": food["cholesterol_mg"]
  253. },
  254. "vitamins": {
  255. "vitamin_a_iu": food["vitamin_a_iu"],
  256. "vitamin_c_mg": food["vitamin_c_mg"]
  257. },
  258. "minerals": {
  259. "calcium_mg": food["calcium_mg"],
  260. "iron_mg": food["iron_mg"],
  261. "potassium_mg": food["potassium_mg"],
  262. "sodium_mg": food["sodium_mg"]
  263. },
  264. "source": food["source"]
  265. }
  266. return structured_data
  267. if __name__ == "__main__":
  268. import uvicorn
  269. uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)