arya-ai-model commited on
Commit
214e325
·
1 Parent(s): 1a96a44

updated model.py

Browse files
Files changed (1) hide show
  1. app.py +9 -3
app.py CHANGED
@@ -1,20 +1,26 @@
1
  import uvicorn
2
  from fastapi import FastAPI, HTTPException
 
3
  from model import generate_code
4
 
5
  app = FastAPI()
6
 
 
 
 
 
 
7
  @app.get("/")
8
  def home():
9
  return {"message": "Code Generation API is running!"}
10
 
11
  @app.post("/generate")
12
- def generate(prompt: str, max_tokens: int = 256):
13
- if not prompt:
14
  raise HTTPException(status_code=400, detail="Prompt cannot be empty.")
15
 
16
  try:
17
- code = generate_code(prompt, max_tokens)
18
  return {"generated_code": code}
19
  except Exception as e:
20
  raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")
 
1
  import uvicorn
2
  from fastapi import FastAPI, HTTPException
3
+ from pydantic import BaseModel
4
  from model import generate_code
5
 
6
  app = FastAPI()
7
 
8
+ # Define request body schema
9
+ class GenerateRequest(BaseModel):
10
+ prompt: str
11
+ max_tokens: int = 256 # Default value
12
+
13
  @app.get("/")
14
  def home():
15
  return {"message": "Code Generation API is running!"}
16
 
17
  @app.post("/generate")
18
+ def generate(request: GenerateRequest): # Expect JSON
19
+ if not request.prompt:
20
  raise HTTPException(status_code=400, detail="Prompt cannot be empty.")
21
 
22
  try:
23
+ code = generate_code(request.prompt, request.max_tokens)
24
  return {"generated_code": code}
25
  except Exception as e:
26
  raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")