from datetime import datetime
from io import IOBase
from typing import List, Optional, Union
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.language_models import BaseLanguageModel
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain_core.prompts import (
ChatPromptTemplate,
MessagesPlaceholder,
)
from langchain_core.prompts.chat import (
BaseMessagePromptTemplate,
HumanMessagePromptTemplate,
)
from langchain_core.tools import BaseTool
from langchain_cohere.chat_models import ChatCohere
from langchain_cohere.csv_agent.prompts import (
CSV_PREAMBLE,
)
# lets define a set of tools for the Agent
from langchain_cohere.csv_agent.tools import (
get_file_peek_tool,
get_file_read_tool,
get_python_tool,
)
[docs]
def create_prompt(
system_message: Optional[BaseMessage] = SystemMessage(
content="You are a helpful AI assistant."
),
extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
) -> ChatPromptTemplate:
"""Create prompt for this agent.
Args:
system_message: Message to use as the system message that will be the
first in the prompt.
extra_prompt_messages: Prompt messages that will be placed between the
system message and the new human input.
Returns:
A prompt template to pass into this agent.
"""
_prompts = extra_prompt_messages or []
messages: List[Union[BaseMessagePromptTemplate, BaseMessage]]
if system_message:
messages = [system_message]
else:
messages = []
messages.extend(
[
*_prompts,
HumanMessagePromptTemplate.from_template("{input}"),
MessagesPlaceholder(variable_name="agent_scratchpad"),
]
)
return ChatPromptTemplate(messages=messages)
def _get_csv_head_str(path: str, number_of_head_rows: int) -> str:
with open(path, "r") as file:
lines = []
for _ in range(number_of_head_rows):
lines.append(file.readline().strip("\n"))
# validate that the head contents are well formatted csv
return " ".join(lines)
def _get_prompt(
path: Union[str, List[str]], number_of_head_rows: int
) -> ChatPromptTemplate:
if isinstance(path, str):
lines = _get_csv_head_str(path, number_of_head_rows)
prompt_message = f"The user uploaded the following attachments:\nFilename: {path}\nWord Count: {count_words_in_file(path)}\nPreview: {lines}" # noqa: E501
elif isinstance(path, list):
prompt_messages = []
for file_path in path:
lines = _get_csv_head_str(file_path, number_of_head_rows)
prompt_messages.append(
f"The user uploaded the following attachments:\nFilename: {file_path}\nWord Count: {count_words_in_file(file_path)}\nPreview: {lines}" # noqa: E501
)
prompt_message = " ".join(prompt_messages)
prompt = create_prompt(system_message=HumanMessage(prompt_message))
return prompt
[docs]
def count_words_in_file(file_path: str) -> int:
try:
with open(file_path, "r") as file:
content = file.readlines()
words = [len(sentence.split()) for sentence in content]
return sum(words)
except FileNotFoundError:
print("File not found.")
return 0
except Exception as e:
print("An error occurred:", str(e))
return 0
[docs]
def create_csv_agent(
llm: BaseLanguageModel,
path: Union[str, List[str]],
extra_tools: List[BaseTool] = [],
pandas_kwargs: Optional[dict] = None,
prompt: Optional[ChatPromptTemplate] = None,
number_of_head_rows: int = 5,
verbose: bool = True,
return_intermediate_steps: bool = True,
temp_path_dir: Optional[str] = None,
temp_path_prefix: Optional[str] = "langchain",
temp_path_suffix: Optional[str] = "csv_agent",
) -> AgentExecutor:
"""Create csv agent with the specified language model.
Args:
llm: Language model to use for the agent.
path: A string path, or a list of string paths
that can be read in as pandas DataFrames with pd.read_csv().
number_of_head_rows: Number of rows to display in the prompt for sample data
include_df_in_prompt: Display the DataFrame sample values in the prompt.
pandas_kwargs: Named arguments to pass to pd.read_csv().
prefix: Prompt prefix string.
suffix: Prompt suffix string.
prompt: Prompt to use for the agent. This takes precedence over the other prompt arguments, such as suffix and prefix.
temp_path_dir: Temporary directory to store the csv files in for the python repl.
delete_temp_path: Whether to delete the temporary directory after the agent is done. This only works if temp_path_dir is not provided.
Returns:
An AgentExecutor with the specified agent_type agent and access to
a PythonREPL and any user-provided extra_tools.
Example:
.. code-block:: python
from langchain_cohere import ChatCohere, create_csv_agent
llm = ChatCohere(model="command-r-plus", temperature=0)
agent_executor = create_csv_agent(
llm,
"titanic.csv"
)
resp = agent_executor.invoke({"input":"How many people were on the titanic?"})
print(resp.get("output"))
""" # noqa: E501
try:
import pandas as pd
except ImportError:
raise ImportError(
"pandas package not found, please install with `pip install pandas`."
)
_kwargs = pandas_kwargs or {}
if isinstance(path, (str)):
df = pd.read_csv(path, **_kwargs)
elif isinstance(path, list):
df = []
for item in path:
if not isinstance(item, (str, IOBase)):
raise ValueError(f"Expected str or file-like object, got {type(path)}")
df.append(pd.read_csv(item, **_kwargs))
else:
raise ValueError(f"Expected str, list, or file-like object, got {type(path)}")
if not prompt:
prompt = _get_prompt(path, number_of_head_rows)
final_tools = [
get_file_read_tool(),
get_file_peek_tool(),
get_python_tool(),
] + extra_tools
if "preamble" in llm.__dict__ and not llm.__dict__.get("preamble"):
llm = ChatCohere(**llm.__dict__)
llm.preamble = CSV_PREAMBLE.format(
current_date=datetime.now().strftime("%A, %B %d, %Y %H:%M:%S")
)
agent = create_tool_calling_agent(llm=llm, tools=final_tools, prompt=prompt)
agent_executor = AgentExecutor(
agent=agent,
tools=final_tools,
verbose=verbose,
return_intermediate_steps=return_intermediate_steps,
)
return agent_executor
