CoolFace
Apppublic

web-agentix/deploy-backend

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
agent.py152 linesDownload Raw Back to root
1import os2import openai3import cohere4from qdrant_client import QdrantClient, models5from dotenv import load_dotenv6import logging7from openai import OpenAI8import json9import sys10 11# Import components from openai-agents SDK12from agents import Agent, Runner, function_tool13 14# Configure logging15logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')16 17# Load environment variables - keep for standalone script use18load_dotenv()19 20# --- Constants ---21EMBEDDING_MODEL = "embed-english-light-v3.0"22QDRANT_COLLECTION_NAME = "reg_embedding"23 24# --- Client Getter Functions (Modified to accept API keys) ---25def get_cohere_client(cohere_api_key: str) -> cohere.Client:26    if not cohere_api_key:27        raise ValueError("COHERE_API_KEY not found. Please set it in your .env file.")28    return cohere.Client(cohere_api_key)29 30def get_qdrant_client(qdrant_url: str, qdrant_api_key: str) -> QdrantClient:31    if not qdrant_url or not qdrant_api_key:32        raise ValueError("QDRANT_URL and QDRANT_API_KEY not found. Please set them in your .env file.")33    return QdrantClient(url=qdrant_url, api_key=qdrant_api_key)34 35def get_openai_client(openai_api_key: str) -> OpenAI:36    if not openai_api_key:37        raise ValueError("OPENAI_API_KEY not found. Please set it in your .env file.")38    return OpenAI(api_key=openai_api_key)39 40# --- Function Tools ---41@function_tool42def retrieve_from_qdrant(query: str, top_k: int = 5, score_threshold: float = 0.0) -> str:43    """44    Retrieves relevant text chunks from Qdrant based on a natural language query.45    This function acts as a tool for the AI agent.46    """47    logging.info(f"Retrieval tool called with query: '{query}', top_k: {top_k}, score_threshold: {score_threshold}")48    49    # Get API keys from environment (assuming they are loaded by the caller)50    COHERE_API_KEY = os.getenv("COHERE_API_KEY")51    QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")52    QDRANT_URL = os.getenv("QDRANT_URL")53 54    try:55        cohere_client_instance = get_cohere_client(COHERE_API_KEY)56        qdrant_client_instance = get_qdrant_client(QDRANT_URL, QDRANT_API_KEY)57 58        query_embedding = cohere_client_instance.embed(59            texts=[query],60            model=EMBEDDING_MODEL,61            input_type="search_query"62        ).embeddings[0]63 64        search_result = qdrant_client_instance.query_points(65            collection_name=QDRANT_COLLECTION_NAME,66            query=query_embedding,67            limit=top_k,68            score_threshold=score_threshold69        )70        71        if not search_result.points:72            logging.info("No relevant results found by Qdrant retrieval tool.")73            return "No relevant information found in the knowledge base."74 75        retrieved_texts = [hit.payload.get('text_snippet', '') for hit in search_result.points]76        77        return "\n\n".join(retrieved_texts)78 79    except Exception as e:80        logging.error(f"Error during Qdrant retrieval: {e}")81        return "An error occurred while retrieving information."82 83# Define the retrieval tool for the OpenAI agent84retrieval_tool = {85    "type": "function",86    "function": {87        "name": "retrieve_from_qdrant",88        "description": "Retrieves relevant text chunks from the knowledge base using a natural language query.",89        "parameters": {90            "type": "object",91            "properties": {92                "query": {93                    "type": "string",94                    "description": "The natural language query for retrieval."95                },96                "top_k": {97                    "type": "integer",98                    "description": "The maximum number of relevant chunks to retrieve.",99                    "default": 5100                },101                "score_threshold": {102                    "type": "number",103                    "description": "The minimum relevance score for retrieved chunks.",104                    "default": 0.0105                }106            },107            "required": ["query"]108        }109    }110}111 112# --- Agent Main Function ---113def main(user_query: str):114    """115    Main function to run the AI agent.116    """117    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # Ensure API key is loaded for agent init118    openai_client_instance = get_openai_client(OPENAI_API_KEY)119 120    agent = Agent(121        name="RoboticsExpert",122        instructions="You are a helpful AI assistant specialized in Physical AI and Humanoid Robotics. Answer questions based on the provided tools and retrieved information. Only use the retrieve_from_qdrant tool to get information from the knowledge base.",123        tools=[retrieve_from_qdrant],124        model="gpt-4-0613",125    )126    127    messages = [128        {"role": "system", "content": "You are a helpful AI assistant specialized in Physical AI and Humanoid Robotics. Answer questions based on the provided tools and retrieved information."},129        {"role": "user", "content": user_query}130    ]131    tools = [retrieval_tool]132    133    try:134        logging.info(f"Agent received query: {user_query}")135        result = Runner.run_sync(agent, user_query)136        print(result.final_output)137 138    except openai.APIError as e:139        logging.error(f"OpenAI API Error: {e}")140        print("An error occurred with the OpenAI API. Please check your API key and network connection.")141    except Exception as e:142        logging.error(f"An unexpected error occurred during agent interaction: {e}")143        print("An unexpected error occurred while processing your request.")144 145 146if __name__ == "__main__":147    if len(sys.argv) > 1:148        user_query = " ".join(sys.argv[1:])149        main(user_query)150    else:151        print("Please provide a query as an argument. Example: python agent.py \"What is ROS?\"")152