Spaces:
Sleeping
Sleeping
File size: 5,795 Bytes
8e1982e 970df01 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
from fastapi import APIRouter, HTTPException, Request
from api.models import PromptRequest, HistoryRequest
from core.llm import llm, GROQ_MODEL
from core.memory import get_memory, conversation_memories
from core.utils import classify_message, process_response
import uuid
import time
from langchain.schema import HumanMessage, AIMessage
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, ConversationChain
router = APIRouter()
# Prompt templates
CONVERSATION_TEMPLATE = """You are a helpful AI assistant. Have a natural conversation with the user.
Current conversation:
{history}
Human: {input}
AI:"""
CODE_GENERATION_TEMPLATE = """You are an expert Python programmer. Generate clean, efficient Python code for the given request.
Request: {input}
Provide only the Python code without explanation:
```python"""
EXPLANATION_TEMPLATE = """You are a programming tutor. Explain how to solve the programming task clearly without providing code.
Task: {input}
Explanation:"""
BOTH_TEMPLATE = """You are an expert Python programmer and tutor. For the given request, provide both a clear explanation and Python code.
Request: {input}
First, provide a clear explanation of the approach:
[EXPLANATION]
Then, provide the Python code:
```python"""
def create_chain(template: str, memory=None):
input_vars = ["input"]
if memory:
input_vars.append("history")
prompt = PromptTemplate(
input_variables=input_vars,
template=template
)
if memory:
return ConversationChain(
llm=llm,
prompt=prompt,
memory=memory,
verbose=True,
output_key="output"
)
else:
return LLMChain(
llm=llm,
prompt=prompt,
verbose=True,
output_key="text"
)
@router.get("/")
async def root():
return {
"status": "ok",
"title": "LangChain Chat Bot",
"version": "2.0.0",
"model": GROQ_MODEL,
"description": "Advanced chatbot using LangChain and Groq models",
"features": ["Conversation Memory", "Code Generation", "Explanations", "Session Management"]
}
@router.post("/generate/")
async def generate_response(request: PromptRequest):
if llm is None:
raise HTTPException(status_code=503, detail="AI model not initialized. Check server logs for errors during startup (e.g., missing API key).")
try:
session_id = request.session_id or str(uuid.uuid4())
memory = get_memory(session_id, request.max_history)
message_type = classify_message(request.prompt)
response_data = {}
if message_type == "conversation" or request.response_type == "conversation":
chain = create_chain(CONVERSATION_TEMPLATE, memory)
response = chain.predict(input=request.prompt)
response_data = {"response": response.strip(), "message_type": "conversation"}
elif request.response_type == "code":
chain = create_chain(CODE_GENERATION_TEMPLATE)
response = chain.run(input=request.prompt)
response_data = process_response(response, "code")
response_data["message_type"] = "code"
elif request.response_type == "explanation":
chain = create_chain(EXPLANATION_TEMPLATE)
response = chain.run(input=request.prompt)
response_data = process_response(response, "explanation")
response_data["message_type"] = "explanation"
else: # "both"
chain = create_chain(BOTH_TEMPLATE)
response = chain.run(input=request.prompt)
response_data = process_response(response, "both")
response_data["message_type"] = "both"
response_data["session_id"] = session_id
print(f"✅ Generated response for session {session_id}")
return response_data
except Exception as e:
print(f"❌ Error in generate_response: {str(e)}")
raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")
@router.post("/clear_history/")
async def clear_history(request: HistoryRequest):
if request.session_id in conversation_memories:
conversation_memories[request.session_id].clear()
return {"status": "success", "message": "Conversation history cleared"}
return {"status": "not_found", "message": "Session ID not found"}
@router.post("/get_history/")
async def get_history(request: HistoryRequest):
if request.session_id in conversation_memories:
memory = conversation_memories[request.session_id]
messages = memory.chat_memory.messages
history = []
for msg in messages:
if isinstance(msg, HumanMessage):
history.append(f"Human: {msg.content}")
elif isinstance(msg, AIMessage):
history.append(f"AI: {msg.content}")
return {"status": "success", "history": history}
return {"status": "not_found", "message": "Session ID not found"}
@router.get("/sessions/")
async def get_active_sessions():
return {
"active_sessions": list(conversation_memories.keys()),
"total_sessions": len(conversation_memories)
}
@router.delete("/sessions/{session_id}")
async def delete_session(session_id: str):
if session_id in conversation_memories:
del conversation_memories[session_id]
return {"status": "success", "message": f"Session {session_id} deleted"}
return {"status": "not_found", "message": "Session ID not found"}
@router.get("/health")
async def health_check():
return {
"status": "healthy",
"model": llm.model if llm else None,
"active_sessions": len(conversation_memories),
"langchain_version": "0.0.350"
} |