-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
47 lines (35 loc) · 1.42 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
from llama_index.llms.ollama import Ollama
from llama_parse import LlamaParse
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, PromptTemplate
from llama_index.core.embeddings import resolve_embed_model
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.agent import ReActAgent
from pydantic import BaseModel
from llama_index.core.output_parsers import PydanticOutputParser
from llama_index.core.query_pipeline import QueryPipeline
from prompts import context
from dotenv import load_dotenv
import os
import ast
## code for the agent
load_dotenv()
llm = Ollama(model="mistral", request_timeout=1000.0)
parser = LlamaParse(result_type="markdown")
file_extractor = {".pdf": parser}
documents = SimpleDirectoryReader("./data", file_extractor=file_extractor).load_data()
embed_model = resolve_embed_model("local:BAAI/bge-m3")
vector_index = VectorStoreIndex.from_documents(documents, embed_model=embed_model)
query_engine = vector_index.as_query_engine(llm=llm)
tools = [
QueryEngineTool(
query_engine=query_engine,
metadata=ToolMetadata(
name="law_helper",
description="this a file gives a overview in Morocco's Constitution of 2011 ",
),
),
]
agent = ReActAgent.from_tools(tools, llm=llm, verbose=True, context=context)
while (prompt := input("Enter a prompt (q to quit): ")) != "q":
res=agent.query(prompt)
print(res)