jarguello76's picture
Update app.py
1a01680 verified
import os
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import InferenceClient
from smolagents import BasicAgent
from crewai_tools import BaseTool
# Use your token from HF secrets
HF_TOKEN = os.environ.get("HF_TOKEN") # Must be set in HF Space secrets
# Define a wrapper that conforms to SmolAgent’s LLM expectations
class HuggingFaceInferenceWrapper:
def __init__(self, client: InferenceClient):
self.client = client
def __call__(self, prompt: str, **kwargs):
# Use text_generation endpoint
response = self.client.text_generation(prompt=prompt, max_new_tokens=512)
return response
# Initialize InferenceClient
client = InferenceClient(
model="Qwen/Qwen1.5-1.8B-Chat", # Or another model you have access to
token=HUGGINGFACE_API_KEY
)
# Wrap the client
llm = HuggingFaceInferenceWrapper(client)
# Define a dummy tool
class DuckDuckGoTool(BaseTool):
name: str = "DuckDuckGo Search"
description: str = "Use this tool to search for real-time facts and current events."
def _run(self, query: str) -> str:
return f"Search results for '{query}' (this is a placeholder)."
# Create the agent
agent = BasicAgent(
role="Helpful AI",
goal="Answer tough questions using search and reasoning",
backstory="An expert at finding information and answering questions.",
tools=[DuckDuckGoTool()],
llm=llm
)
# Define FastAPI app
app = FastAPI()
class QuestionInput(BaseModel):
question: str
@app.post("/ask")
def ask_question(payload: QuestionInput):
result = agent.run(payload.question)
return {"answer": result}