Spaces:
Sleeping
Sleeping
| from flask import Flask, request, jsonify | |
| from models import ( | |
| load_llm, text_splitter, embedding_model, prompt_template, | |
| extract_json_block, save_model_state, load_model_state | |
| ) | |
| from langchain.vectorstores import Chroma | |
| app = Flask(__name__) | |
| # Load the model once when the server starts | |
| print("Loading model...") | |
| llm = load_llm() | |
| print("Model loaded successfully!") | |
| def health_check(): | |
| return jsonify({"status": "healthy", "model_loaded": True}) | |
| def generate_workflow(): | |
| try: | |
| data = request.get_json() | |
| user_prompt = data.get('prompt', '') | |
| use_rag = data.get('use_rag', False) | |
| raw_docs = data.get('raw_docs', None) | |
| if not user_prompt.strip(): | |
| return jsonify({"error": "Prompt is required"}), 400 | |
| context_block = "" | |
| if use_rag and raw_docs: | |
| # Process documents for RAG | |
| all_chunks = [] | |
| for doc in raw_docs: | |
| all_chunks.extend(text_splitter.split_text(doc)) | |
| docs_for_vector = [{"page_content": c} for c in all_chunks] | |
| vector = Chroma.from_documents(docs_for_vector, embedding=embedding_model) | |
| retriever = vector.as_retriever(search_kwargs={"k": 3}) | |
| retrieved_docs = retriever.get_relevant_documents(user_prompt) | |
| context_block = "Context:\n" + "\n".join([d.page_content for d in retrieved_docs]) + "\n" | |
| # Generate workflow | |
| full_prompt = prompt_template.format(prompt=user_prompt, context_block=context_block) | |
| wrapped_prompt = f"[INST] {full_prompt} [/INST]" | |
| output = llm(wrapped_prompt, max_tokens=1024, stop=["</s>"]) | |
| generated_text = output["choices"][0]["text"] | |
| generated_json = extract_json_block(generated_text) | |
| if generated_json: | |
| # Optionally save model state for optimization | |
| # save_model_state(llm, "model_state.bin") | |
| return jsonify({ | |
| "success": True, | |
| "workflow": generated_json, | |
| "context_used": bool(context_block) | |
| }) | |
| else: | |
| return jsonify({"error": "Failed to extract valid JSON from generated text"}), 400 | |
| except Exception as e: | |
| return jsonify({"error": str(e)}), 500 | |
| if __name__ == '__main__': | |
| app.run(debug=True, host='0.0.0.0', port=5000) | |