Merge pull request #131 from Madhuvod/legal-changess

Fixed issues in AI Legal Agent
This commit is contained in:
Shubham Saboo 2025-02-25 15:33:16 -06:00 committed by GitHub
commit 81b75cf72c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -7,8 +7,8 @@ from agno.models.openai import OpenAIChat
from agno.embedder.openai import OpenAIEmbedder from agno.embedder.openai import OpenAIEmbedder
import tempfile import tempfile
import os import os
from agno.document.chunking.document import DocumentChunking
#initializing the session state variables
def init_session_state(): def init_session_state():
"""Initialize session state variables""" """Initialize session state variables"""
if 'openai_api_key' not in st.session_state: if 'openai_api_key' not in st.session_state:
@ -23,54 +23,86 @@ def init_session_state():
st.session_state.legal_team = None st.session_state.legal_team = None
if 'knowledge_base' not in st.session_state: if 'knowledge_base' not in st.session_state:
st.session_state.knowledge_base = None st.session_state.knowledge_base = None
# Add a new state variable to track processed files
if 'processed_files' not in st.session_state:
st.session_state.processed_files = set()
COLLECTION_NAME = "legal_documents" # Define your collection name
def init_qdrant(): def init_qdrant():
"""Initialize Qdrant vector database""" """Initialize Qdrant client with configured settings."""
if not st.session_state.qdrant_api_key: if not all([st.session_state.qdrant_api_key, st.session_state.qdrant_url]):
raise ValueError("Qdrant API key not provided") return None
if not st.session_state.qdrant_url: try:
raise ValueError("Qdrant URL not provided") # Create Agno's Qdrant instance which implements VectorDb
vector_db = Qdrant(
return Qdrant( collection=COLLECTION_NAME,
collection="legal_knowledge",
url=st.session_state.qdrant_url, url=st.session_state.qdrant_url,
api_key=st.session_state.qdrant_api_key, api_key=st.session_state.qdrant_api_key,
https=True, embedder=OpenAIEmbedder(
timeout=None, id="text-embedding-3-small",
distance="cosine" api_key=st.session_state.openai_api_key
) )
)
return vector_db
except Exception as e:
st.error(f"🔴 Qdrant connection failed: {str(e)}")
return None
def process_document(uploaded_file, vector_db: Qdrant): def process_document(uploaded_file, vector_db: Qdrant):
"""Process document, create embeddings and store in Qdrant vector database""" """
Process document, create embeddings and store in Qdrant vector database
Args:
uploaded_file: Streamlit uploaded file object
vector_db (Qdrant): Initialized Qdrant instance from Agno
Returns:
PDFKnowledgeBase: Initialized knowledge base with processed documents
"""
if not st.session_state.openai_api_key: if not st.session_state.openai_api_key:
raise ValueError("OpenAI API key not provided") raise ValueError("OpenAI API key not provided")
os.environ['OPENAI_API_KEY'] = st.session_state.openai_api_key os.environ['OPENAI_API_KEY'] = st.session_state.openai_api_key
with tempfile.TemporaryDirectory() as temp_dir:
temp_file_path = os.path.join(temp_dir, uploaded_file.name)
with open(temp_file_path, "wb") as f:
f.write(uploaded_file.getbuffer())
try: try:
# Save the uploaded file to a temporary location
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
temp_file.write(uploaded_file.getvalue())
temp_file_path = temp_file.name
embedder = OpenAIEmbedder( st.info("Loading and processing document...")
model="text-embedding-3-small",
api_key=st.session_state.openai_api_key
)
# Creating knowledge base with explicit Qdrant configuration # Create a PDFKnowledgeBase with the vector_db
knowledge_base = PDFKnowledgeBase( knowledge_base = PDFKnowledgeBase(
path=temp_dir, path=temp_file_path, # Single string path, not a list
vector_db=vector_db, vector_db=vector_db,
reader=PDFReader(chunk=True), reader=PDFReader(),
embedder=embedder, chunking_strategy=DocumentChunking(
recreate_vector_db=True chunk_size=1000,
overlap=200
) )
knowledge_base.load() )
return knowledge_base
# Load the documents into the knowledge base
with st.spinner('📤 Loading documents into knowledge base...'):
try:
knowledge_base.load(recreate=True, upsert=True)
st.success("✅ Documents stored successfully!")
except Exception as e: except Exception as e:
st.error(f"Error loading documents: {str(e)}")
raise
# Clean up the temporary file
try:
os.unlink(temp_file_path)
except Exception:
pass
return knowledge_base
except Exception as e:
st.error(f"Document processing error: {str(e)}")
raise Exception(f"Error processing document: {str(e)}") raise Exception(f"Error processing document: {str(e)}")
def main(): def main():
@ -102,7 +134,7 @@ def main():
qdrant_url = st.text_input( qdrant_url = st.text_input(
"Qdrant URL", "Qdrant URL",
value=st.session_state.qdrant_url if st.session_state.qdrant_url else "https://f499085c-b4bf-4bda-a9a5-227f62a9ca20.us-west-2-0.aws.cloud.qdrant.io:6333", value=st.session_state.qdrant_url if st.session_state.qdrant_url else "",
help="Enter your Qdrant instance URL" help="Enter your Qdrant instance URL"
) )
if qdrant_url: if qdrant_url:
@ -111,7 +143,9 @@ def main():
if all([st.session_state.qdrant_api_key, st.session_state.qdrant_url]): if all([st.session_state.qdrant_api_key, st.session_state.qdrant_url]):
try: try:
if not st.session_state.vector_db: if not st.session_state.vector_db:
# Make sure we're initializing a QdrantClient here
st.session_state.vector_db = init_qdrant() st.session_state.vector_db = init_qdrant()
if st.session_state.vector_db:
st.success("Successfully connected to Qdrant!") st.success("Successfully connected to Qdrant!")
except Exception as e: except Exception as e:
st.error(f"Failed to connect to Qdrant: {str(e)}") st.error(f"Failed to connect to Qdrant: {str(e)}")
@ -123,16 +157,23 @@ def main():
uploaded_file = st.file_uploader("Upload Legal Document", type=['pdf']) uploaded_file = st.file_uploader("Upload Legal Document", type=['pdf'])
if uploaded_file: if uploaded_file:
# Check if this file has already been processed
if uploaded_file.name not in st.session_state.processed_files:
with st.spinner("Processing document..."): with st.spinner("Processing document..."):
try: try:
# Process the document and get the knowledge base
knowledge_base = process_document(uploaded_file, st.session_state.vector_db) knowledge_base = process_document(uploaded_file, st.session_state.vector_db)
if knowledge_base:
st.session_state.knowledge_base = knowledge_base st.session_state.knowledge_base = knowledge_base
# Add the file to processed files
st.session_state.processed_files.add(uploaded_file.name)
# Initialize agents # Initialize agents
legal_researcher = Agent( legal_researcher = Agent(
name="Legal Researcher", name="Legal Researcher",
role="Legal research specialist", role="Legal research specialist",
model=OpenAIChat(model="gpt-4o"), model=OpenAIChat(id="gpt-4o"),
tools=[DuckDuckGoTools()], tools=[DuckDuckGoTools()],
knowledge=st.session_state.knowledge_base, knowledge=st.session_state.knowledge_base,
search_knowledge=True, search_knowledge=True,
@ -149,8 +190,8 @@ def main():
contract_analyst = Agent( contract_analyst = Agent(
name="Contract Analyst", name="Contract Analyst",
role="Contract analysis specialist", role="Contract analysis specialist",
model=OpenAIChat(model="gpt-4o"), model=OpenAIChat(id="gpt-4o"),
knowledge=knowledge_base, knowledge=st.session_state.knowledge_base,
search_knowledge=True, search_knowledge=True,
instructions=[ instructions=[
"Review contracts thoroughly", "Review contracts thoroughly",
@ -163,8 +204,8 @@ def main():
legal_strategist = Agent( legal_strategist = Agent(
name="Legal Strategist", name="Legal Strategist",
role="Legal strategy specialist", role="Legal strategy specialist",
model=OpenAIChat(model="gpt-4o"), model=OpenAIChat(id="gpt-4o"),
knowledge=knowledge_base, knowledge=st.session_state.knowledge_base,
search_knowledge=True, search_knowledge=True,
instructions=[ instructions=[
"Develop comprehensive legal strategies", "Develop comprehensive legal strategies",
@ -178,7 +219,7 @@ def main():
st.session_state.legal_team = Agent( st.session_state.legal_team = Agent(
name="Legal Team Lead", name="Legal Team Lead",
role="Legal team coordinator", role="Legal team coordinator",
model=OpenAIChat(model="gpt-4o"), model=OpenAIChat(id="gpt-4o"),
team=[legal_researcher, contract_analyst, legal_strategist], team=[legal_researcher, contract_analyst, legal_strategist],
knowledge=st.session_state.knowledge_base, knowledge=st.session_state.knowledge_base,
search_knowledge=True, search_knowledge=True,
@ -197,6 +238,9 @@ def main():
except Exception as e: except Exception as e:
st.error(f"Error processing document: {str(e)}") st.error(f"Error processing document: {str(e)}")
else:
# File already processed, just show a message
st.success("✅ Document already processed and team ready!")
st.divider() st.divider()
st.header("🔍 Analysis Options") st.header("🔍 Analysis Options")