You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
103 lines
2.6 KiB
103 lines
2.6 KiB
from abc import ABC
|
|
from typing import Any, List, Tuple
|
|
|
|
from deepsearcher.vector_db import RetrievalResult
|
|
|
|
|
|
def describe_class(description):
|
|
"""
|
|
Decorator function to add a description to a class.
|
|
|
|
This decorator adds a __description__ attribute to the decorated class,
|
|
which can be used for documentation or introspection.
|
|
|
|
Args:
|
|
description: The description to add to the class.
|
|
|
|
Returns:
|
|
A decorator function that adds the description to the class.
|
|
"""
|
|
|
|
def decorator(cls):
|
|
cls.__description__ = description
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
|
|
class BaseAgent(ABC):
|
|
"""
|
|
Abstract base class for all agents in the DeepSearcher system.
|
|
|
|
This class defines the basic interface for agents, including initialization
|
|
and invocation methods.
|
|
"""
|
|
|
|
def __init__(self, **kwargs):
|
|
"""
|
|
Initialize a BaseAgent object.
|
|
|
|
Args:
|
|
**kwargs: Arbitrary keyword arguments.
|
|
"""
|
|
pass
|
|
|
|
def invoke(self, query: str, **kwargs) -> Any:
|
|
"""
|
|
Invoke the agent and return the result.
|
|
|
|
Args:
|
|
query: The query string.
|
|
**kwargs: Additional keyword arguments.
|
|
|
|
Returns:
|
|
The result of invoking the agent.
|
|
"""
|
|
|
|
|
|
class RAGAgent(BaseAgent):
|
|
"""
|
|
Abstract base class for Retrieval-Augmented Generation (RAG) agents.
|
|
|
|
This class extends BaseAgent with methods specific to RAG, including
|
|
retrieval and query methods.
|
|
"""
|
|
|
|
def __init__(self, **kwargs):
|
|
"""
|
|
Initialize a RAGAgent object.
|
|
|
|
Args:
|
|
**kwargs: Arbitrary keyword arguments.
|
|
"""
|
|
pass
|
|
|
|
def retrieve(self, query: str, **kwargs) -> Tuple[List[RetrievalResult], int, dict]:
|
|
"""
|
|
Retrieve document results from the knowledge base.
|
|
|
|
Args:
|
|
query: The query string.
|
|
**kwargs: Additional keyword arguments.
|
|
|
|
Returns:
|
|
A tuple containing:
|
|
- the retrieved results
|
|
- the total number of token usages of the LLM
|
|
- any additional metadata, which can be an empty dictionary
|
|
"""
|
|
|
|
def query(self, query: str, **kwargs) -> Tuple[str, List[RetrievalResult], int]:
|
|
"""
|
|
Query the agent and return the answer.
|
|
|
|
Args:
|
|
query: The query string.
|
|
**kwargs: Additional keyword arguments.
|
|
|
|
Returns:
|
|
A tuple containing:
|
|
- the result generated from LLM
|
|
- the retrieved document results
|
|
- the total number of token usages of the LLM
|
|
"""
|
|
|