chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,164 @@
|
||||
import pandas as pd
|
||||
from typing import List, Dict
|
||||
from autoagent.memory.rag_memory import Memory, Reranker
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from litellm import completion
|
||||
from pydantic import BaseModel
|
||||
"""
|
||||
Category | Tool_Name | Tool_Description | API_Name | API_Description | Method | API_Details | Required_API_Key | Platform
|
||||
"""
|
||||
class ToolMemory(Memory):
|
||||
def __init__(
|
||||
self,
|
||||
project_path: str,
|
||||
db_name: str = '.tool_table',
|
||||
platform: str = 'OpenAI',
|
||||
api_key: str = None,
|
||||
embedding_model: str = "text-embedding-3-small",
|
||||
):
|
||||
super().__init__(
|
||||
project_path=project_path,
|
||||
db_name=db_name,
|
||||
platform=platform,
|
||||
api_key=api_key,
|
||||
embedding_model=embedding_model
|
||||
)
|
||||
self.collection_name = 'tool_memory'
|
||||
|
||||
def add_dataframe(self, df: pd.DataFrame, collection: str = None, batch_size: int = 100):
|
||||
if not collection:
|
||||
collection = self.collection_name
|
||||
queries = []
|
||||
for idx, row in df.iterrows():
|
||||
query = {
|
||||
'query': ' '.join(row[['Tool_Name', 'Tool_Description', 'API_Name', 'API_Description']].astype(str)),
|
||||
'response': row.to_json()
|
||||
}
|
||||
queries.append(query)
|
||||
|
||||
# self.add_query(queries, collection=collection)
|
||||
print(f'Adding {len(queries)} queries to {collection} with batch size {batch_size}')
|
||||
num_batches = math.ceil(len(queries) / batch_size)
|
||||
|
||||
for i in range(num_batches):
|
||||
start_idx = i * batch_size
|
||||
end_idx = min((i + 1) * batch_size, len(queries))
|
||||
batch_queries = queries[start_idx:end_idx]
|
||||
|
||||
# Add the current batch of queries
|
||||
self.add_query(batch_queries, collection=collection)
|
||||
|
||||
print(f"Batch {i+1}/{num_batches} added")
|
||||
|
||||
def query_table(
|
||||
self,
|
||||
query_text: str,
|
||||
collection: str = None,
|
||||
n_results: int = 5
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Query the table and return the results
|
||||
"""
|
||||
if not collection:
|
||||
collection = self.collection_name
|
||||
results = self.query([query_text], collection=collection, n_results=n_results)
|
||||
|
||||
metadata_results = results['metadatas'][0]
|
||||
|
||||
df_results = pd.DataFrame([json.loads(item['response']) for item in metadata_results])
|
||||
return df_results
|
||||
|
||||
def peek_table(self, collection: str = None, n_results: int = 20) -> pd.DataFrame:
|
||||
"""
|
||||
Peek at the data in the table
|
||||
"""
|
||||
if not collection:
|
||||
collection = self.collection_name
|
||||
results = self.peek(collection=collection, n_results=n_results)
|
||||
df_results = pd.DataFrame([json.loads(item['response']) for item in results['metadatas']])
|
||||
return df_results
|
||||
|
||||
class ToolReranker(Reranker):
|
||||
def rerank(self, query_text: str, query_df: pd.DataFrame) -> str:
|
||||
system_prompt = \
|
||||
"""
|
||||
You are a helpful assistant that reranks the given API table based on the query.
|
||||
You should select the top 5 APIs to answer the query in the given format.
|
||||
You can only select APIs I give you.
|
||||
Directly give the answer without any other words.
|
||||
"""
|
||||
# Use the DataFrame's to_dict method to convert all rows to a list of dictionaries
|
||||
# print('query_df', query_df)
|
||||
api_data = query_df.to_dict(orient='records')
|
||||
|
||||
# Use a list comprehension and f-string to format each API's data
|
||||
api_prompts = [f"\n\nAPI {i+1}:\n{api}" for i, api in enumerate(api_data)]
|
||||
|
||||
# add the query text to the prompt
|
||||
prompt = ''.join(api_prompts)
|
||||
prompt = f"The query is: {query_text}\n\n{prompt}"
|
||||
message = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": prompt}
|
||||
]
|
||||
|
||||
class Tools(BaseModel):
|
||||
tool_name: str
|
||||
api_name: str
|
||||
rank: int
|
||||
|
||||
class RerankResult(BaseModel):
|
||||
tools: list[Tools]
|
||||
|
||||
create_params = {
|
||||
"model": self.model,
|
||||
"messages": message,
|
||||
"stream": False,
|
||||
"response_format": RerankResult
|
||||
}
|
||||
response = completion(**create_params).choices[0].message.content
|
||||
print(response)
|
||||
rerank_result = json.loads(response)["tools"]
|
||||
print(rerank_result)
|
||||
if len(rerank_result) == 0:
|
||||
return "Fail to retrieve the relevant information from the tool documentation."
|
||||
try:
|
||||
return self.wrap_rerank_result(rerank_result, query_df)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to wrap rerank result: {e}")
|
||||
|
||||
def wrap_rerank_result(self, rerank_result: List[pd.DataFrame], query_df: pd.DataFrame) -> str:
|
||||
res = ""
|
||||
res_tmp = """
|
||||
The rank {rank} referenced tool documentation is:
|
||||
API Name: {api_name}
|
||||
API Description: {api_description}
|
||||
API Details: {api_details}
|
||||
Required API Key: {required_api_key}
|
||||
Platform: {platform}
|
||||
"""
|
||||
try:
|
||||
for tool_api in rerank_result:
|
||||
tool_name = tool_api['tool_name']
|
||||
api_name = tool_api['api_name']
|
||||
matched_rows = query_df[(query_df['API_Name'] == api_name) & (query_df['Tool_Name'] == tool_name)]
|
||||
if not matched_rows.empty:
|
||||
res = res + res_tmp.format(rank=tool_api['rank'], api_name=matched_rows['API_Name'].values[0], api_description=matched_rows['API_Description'].values[0], api_details=matched_rows['API_Details'].values[0], required_api_key=matched_rows['Required_API_Key'].values[0], platform=matched_rows['Platform'].values[0])
|
||||
return res
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to wrap rerank result: {e}")
|
||||
def dummy_rerank(self, query_text: str, query_df: pd.DataFrame) -> str:
|
||||
res = ""
|
||||
res_tmp = """
|
||||
The rank {rank} referenced tool documentation is:
|
||||
API Name: {api_name}
|
||||
API Description: {api_description}
|
||||
API Details: {api_details}
|
||||
Required API Key: {required_api_key}
|
||||
Platform: {platform}
|
||||
"""
|
||||
for i in range(len(query_df)):
|
||||
res = res + res_tmp.format(rank=i+1, api_name=query_df['API_Name'].values[i], api_description=query_df['API_Description'].values[i], api_details=query_df['API_Details'].values[i], required_api_key=query_df['Required_API_Key'].values[i], platform=query_df['Platform'].values[i])
|
||||
return res
|
||||
Reference in New Issue
Block a user