Your Name
2 months ago
parent
9d471c2188
commit
c1c98a85af
@ -0,0 +1,54 @@
|
||||
import pandas as pd
|
||||
import json
|
||||
from loguru import logger
|
||||
|
||||
|
||||
def dict_to_dataframe(data: dict) -> pd.DataFrame:
|
||||
"""
|
||||
Converts a dictionary into a Pandas DataFrame with formatted values.
|
||||
Handles non-serializable values gracefully by skipping them.
|
||||
|
||||
Args:
|
||||
data (dict): The dictionary to convert.
|
||||
|
||||
Returns:
|
||||
pd.DataFrame: A DataFrame representation of the dictionary.
|
||||
"""
|
||||
formatted_data = {}
|
||||
|
||||
for key, value in data.items():
|
||||
try:
|
||||
# Attempt to serialize the value
|
||||
if isinstance(value, list):
|
||||
# Format list as comma-separated string
|
||||
formatted_value = ", ".join(
|
||||
str(item) for item in value
|
||||
)
|
||||
elif isinstance(value, dict):
|
||||
# Format dict as key-value pairs
|
||||
formatted_value = ", ".join(
|
||||
f"{k}: {v}" for k, v in value.items()
|
||||
)
|
||||
else:
|
||||
# Convert other serializable types to string
|
||||
formatted_value = json.dumps(
|
||||
value
|
||||
) # Serialize value to string
|
||||
|
||||
formatted_data[key] = formatted_value
|
||||
except (TypeError, ValueError) as e:
|
||||
# Log and skip non-serializable items
|
||||
logger.warning(
|
||||
f"Skipping non-serializable key '{key}': {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Convert the formatted dictionary into a DataFrame
|
||||
return pd.DataFrame(
|
||||
list(formatted_data.items()), columns=["Key", "Value"]
|
||||
)
|
||||
|
||||
|
||||
example = dict_to_dataframe(data={"chicken": "noodle_soup"})
|
||||
# formatter.print_panel(example)
|
||||
print(example)
|
||||
@ -0,0 +1,228 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
|
||||
from loguru import logger
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
from swarms import Agent, AgentRearrange
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Get the OpenAI API key from the environment variable
|
||||
api_key = os.getenv("GROQ_API_KEY")
|
||||
|
||||
# Model
|
||||
model = OpenAIChat(
|
||||
openai_api_base="https://api.groq.com/openai/v1",
|
||||
openai_api_key=api_key,
|
||||
model_name="llama-3.1-70b-versatile",
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
|
||||
class LlamaIndexDB:
|
||||
"""A class to manage document indexing and querying using LlamaIndex.
|
||||
|
||||
This class provides functionality to add documents from a directory and query the indexed documents.
|
||||
|
||||
Args:
|
||||
data_dir (str): Directory containing documents to index. Defaults to "docs".
|
||||
**kwargs: Additional arguments passed to SimpleDirectoryReader and VectorStoreIndex.
|
||||
SimpleDirectoryReader kwargs:
|
||||
- filename_as_id (bool): Use filenames as document IDs
|
||||
- recursive (bool): Recursively read subdirectories
|
||||
- required_exts (List[str]): Only read files with these extensions
|
||||
- exclude_hidden (bool): Skip hidden files
|
||||
|
||||
VectorStoreIndex kwargs:
|
||||
- service_context: Custom service context
|
||||
- embed_model: Custom embedding model
|
||||
- similarity_top_k (int): Number of similar docs to retrieve
|
||||
- store_nodes_override (bool): Override node storage
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str = "docs", **kwargs) -> None:
|
||||
"""Initialize the LlamaIndexDB with an empty index.
|
||||
|
||||
Args:
|
||||
data_dir (str): Directory containing documents to index
|
||||
**kwargs: Additional arguments for SimpleDirectoryReader and VectorStoreIndex
|
||||
"""
|
||||
self.data_dir = data_dir
|
||||
self.index: Optional[VectorStoreIndex] = None
|
||||
self.reader_kwargs = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k
|
||||
in SimpleDirectoryReader.__init__.__code__.co_varnames
|
||||
}
|
||||
self.index_kwargs = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k not in self.reader_kwargs
|
||||
}
|
||||
|
||||
logger.info("Initialized LlamaIndexDB")
|
||||
data_path = Path(self.data_dir)
|
||||
if not data_path.exists():
|
||||
logger.error(f"Directory not found: {self.data_dir}")
|
||||
raise FileNotFoundError(
|
||||
f"Directory {self.data_dir} does not exist"
|
||||
)
|
||||
|
||||
try:
|
||||
documents = SimpleDirectoryReader(
|
||||
self.data_dir, **self.reader_kwargs
|
||||
).load_data()
|
||||
self.index = VectorStoreIndex.from_documents(
|
||||
documents, **self.index_kwargs
|
||||
)
|
||||
logger.success(
|
||||
f"Successfully indexed documents from {self.data_dir}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error indexing documents: {str(e)}")
|
||||
raise
|
||||
|
||||
def query(self, query: str, **kwargs) -> str:
|
||||
"""Query the indexed documents.
|
||||
|
||||
Args:
|
||||
query (str): The query string to search for
|
||||
**kwargs: Additional arguments passed to the query engine
|
||||
- similarity_top_k (int): Number of similar documents to retrieve
|
||||
- streaming (bool): Enable streaming response
|
||||
- response_mode (str): Response synthesis mode
|
||||
- max_tokens (int): Maximum tokens in response
|
||||
|
||||
Returns:
|
||||
str: The response from the query engine
|
||||
|
||||
Raises:
|
||||
ValueError: If no documents have been indexed yet
|
||||
"""
|
||||
if self.index is None:
|
||||
logger.error("No documents have been indexed yet")
|
||||
raise ValueError("Must add documents before querying")
|
||||
|
||||
try:
|
||||
query_engine = self.index.as_query_engine(**kwargs)
|
||||
response = query_engine.query(query)
|
||||
print(response)
|
||||
logger.info(f"Successfully queried: {query}")
|
||||
return str(response)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during query: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Initialize specialized medical agents
|
||||
medical_data_extractor = Agent(
|
||||
agent_name="Medical-Data-Extractor",
|
||||
system_prompt="You are a specialized medical data extraction expert, trained in processing and analyzing clinical data, lab results, medical imaging reports, and patient records. Your role is to carefully extract relevant medical information while maintaining strict HIPAA compliance and patient confidentiality. Focus on identifying key clinical indicators, test results, vital signs, medication histories, and relevant patient history. Pay special attention to temporal relationships between symptoms, treatments, and outcomes. Ensure all extracted data maintains proper medical context and terminology.",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="medical_data_extractor.json",
|
||||
user_name="medical_team",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
diagnostic_specialist = Agent(
|
||||
agent_name="Diagnostic-Specialist",
|
||||
system_prompt="You are a senior diagnostic physician with extensive experience in differential diagnosis. Your role is to analyze patient symptoms, lab results, and clinical findings to develop comprehensive diagnostic assessments. Consider all presenting symptoms, patient history, risk factors, and test results to formulate possible diagnoses. Prioritize diagnoses based on clinical probability and severity. Always consider both common and rare conditions that match the symptom pattern. Recommend additional tests or imaging when needed for diagnostic clarity. Follow evidence-based diagnostic criteria and current medical guidelines.",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="diagnostic_specialist.json",
|
||||
user_name="medical_team",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
treatment_planner = Agent(
|
||||
agent_name="Treatment-Planner",
|
||||
system_prompt="You are an experienced clinical treatment specialist focused on developing comprehensive treatment plans. Your expertise covers both acute and chronic condition management, medication selection, and therapeutic interventions. Consider patient-specific factors including age, comorbidities, allergies, and contraindications when recommending treatments. Incorporate both pharmacological and non-pharmacological interventions. Emphasize evidence-based treatment protocols while considering patient preferences and quality of life. Address potential drug interactions and side effects. Include monitoring parameters and treatment milestones.",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="treatment_planner.json",
|
||||
user_name="medical_team",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
specialist_consultant = Agent(
|
||||
agent_name="Specialist-Consultant",
|
||||
system_prompt="You are a medical specialist consultant with expertise across multiple disciplines including cardiology, neurology, endocrinology, and internal medicine. Your role is to provide specialized insight for complex cases requiring deep domain knowledge. Analyze cases from your specialist perspective, considering rare conditions and complex interactions between multiple systems. Provide detailed recommendations for specialized testing, imaging, or interventions within your domain. Highlight potential complications or considerations that may not be immediately apparent to general practitioners.",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="specialist_consultant.json",
|
||||
user_name="medical_team",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
patient_care_coordinator = Agent(
|
||||
agent_name="Patient-Care-Coordinator",
|
||||
system_prompt="You are a patient care coordinator specializing in comprehensive healthcare management. Your role is to ensure holistic patient care by coordinating between different medical specialists, considering patient needs, and managing care transitions. Focus on patient education, medication adherence, lifestyle modifications, and follow-up care planning. Consider social determinants of health, patient resources, and access to care. Develop actionable care plans that patients can realistically follow. Coordinate with other healthcare providers to ensure continuity of care and proper implementation of treatment plans.",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="patient_care_coordinator.json",
|
||||
user_name="medical_team",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
|
||||
# Initialize the SwarmRouter to coordinate the medical agents
|
||||
router = AgentRearrange(
|
||||
name="medical-diagnosis-treatment-swarm",
|
||||
description="Collaborative medical team for comprehensive patient diagnosis and treatment planning",
|
||||
max_loops=1, # Limit to one iteration through the agent flow
|
||||
agents=[
|
||||
medical_data_extractor, # First agent to extract medical data
|
||||
diagnostic_specialist, # Second agent to analyze and diagnose
|
||||
treatment_planner, # Third agent to plan treatment
|
||||
specialist_consultant, # Fourth agent to provide specialist input
|
||||
patient_care_coordinator, # Final agent to coordinate care plan
|
||||
],
|
||||
# Configure the document storage and retrieval system
|
||||
memory_system=LlamaIndexDB(
|
||||
data_dir="docs", # Directory containing medical documents
|
||||
filename_as_id=True, # Use filenames as document identifiers
|
||||
recursive=True, # Search subdirectories
|
||||
# required_exts=[".txt", ".pdf", ".docx"], # Supported file types
|
||||
similarity_top_k=10, # Return top 10 most relevant documents
|
||||
),
|
||||
# Define the sequential flow of information between agents
|
||||
flow=f"{medical_data_extractor.agent_name} -> {diagnostic_specialist.agent_name} -> {treatment_planner.agent_name} -> {specialist_consultant.agent_name} -> {patient_care_coordinator.agent_name}",
|
||||
)
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
# Run a comprehensive medical analysis task for patient Lucas Brown
|
||||
router.run(
|
||||
"Analyze this Lucas Brown's medical data to provide a diagnosis and treatment plan"
|
||||
)
|
||||
@ -0,0 +1,63 @@
|
||||
import os
|
||||
import google.generativeai as genai
|
||||
from loguru import logger
|
||||
|
||||
|
||||
class GeminiModel:
|
||||
"""
|
||||
Represents a GeminiModel instance for generating text based on user input.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
top_k: float,
|
||||
):
|
||||
"""
|
||||
Initializes the GeminiModel by setting up the API key, generation configuration, and starting a chat session.
|
||||
Raises a KeyError if the GEMINI_API_KEY environment variable is not found.
|
||||
"""
|
||||
try:
|
||||
api_key = os.environ["GEMINI_API_KEY"]
|
||||
genai.configure(api_key=api_key)
|
||||
self.generation_config = {
|
||||
"temperature": 1,
|
||||
"top_p": 0.95,
|
||||
"top_k": 40,
|
||||
"max_output_tokens": 8192,
|
||||
"response_mime_type": "text/plain",
|
||||
}
|
||||
self.model = genai.GenerativeModel(
|
||||
model_name="gemini-1.5-pro",
|
||||
generation_config=self.generation_config,
|
||||
)
|
||||
self.chat_session = self.model.start_chat(history=[])
|
||||
except KeyError as e:
|
||||
logger.error(f"Environment variable not found: {e}")
|
||||
raise
|
||||
|
||||
def run(self, task: str) -> str:
|
||||
"""
|
||||
Sends a message to the chat session and returns the response text.
|
||||
Raises an Exception if there's an error running the GeminiModel.
|
||||
|
||||
Args:
|
||||
task (str): The input task or message to send to the chat session.
|
||||
|
||||
Returns:
|
||||
str: The response text from the chat session.
|
||||
"""
|
||||
try:
|
||||
response = self.chat_session.send_message(task)
|
||||
return response.text
|
||||
except Exception as e:
|
||||
logger.error(f"Error running GeminiModel: {e}")
|
||||
raise
|
||||
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
gemini_model = GeminiModel()
|
||||
output = gemini_model.run("INSERT_INPUT_HERE")
|
||||
print(output)
|
||||
@ -0,0 +1,238 @@
|
||||
"""
|
||||
Todo
|
||||
|
||||
- You send structured data to the swarm through the users form they make
|
||||
- then connect rag for every agent using llama index to remember all the students data
|
||||
- structured outputs
|
||||
"""
|
||||
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from swarms import Agent, AgentRearrange
|
||||
from swarm_models import OpenAIChat, OpenAIFunctionCaller
|
||||
from pydantic import BaseModel
|
||||
from typing import List
|
||||
|
||||
|
||||
class CollegeLog(BaseModel):
|
||||
college_name: str
|
||||
college_description: str
|
||||
college_admission_requirements: str
|
||||
|
||||
|
||||
class CollegesRecommendation(BaseModel):
|
||||
colleges: List[CollegeLog]
|
||||
reasoning: str
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Get the API key from environment variable
|
||||
api_key = os.getenv("GROQ_API_KEY")
|
||||
|
||||
# Initialize the model
|
||||
model = OpenAIChat(
|
||||
openai_api_base="https://api.groq.com/openai/v1",
|
||||
openai_api_key=api_key,
|
||||
model_name="llama-3.1-70b-versatile",
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
FINAL_AGENT_PROMPT = """
|
||||
You are a college selection final decision maker. Your role is to:
|
||||
1. Synthesize all previous analyses and discussions
|
||||
2. Weigh competing factors and trade-offs
|
||||
3. Create a final ranked list of recommended colleges
|
||||
4. Provide clear rationale for each recommendation
|
||||
5. Include specific action items for each selected school
|
||||
6. Outline next steps in the application process
|
||||
|
||||
Focus on creating actionable, well-reasoned final recommendations that
|
||||
balance all relevant factors and stakeholder input.
|
||||
|
||||
"""
|
||||
|
||||
function_caller = OpenAIFunctionCaller(
|
||||
system_prompt=FINAL_AGENT_PROMPT,
|
||||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_model=CollegesRecommendation,
|
||||
parallel_tool_calls=True,
|
||||
)
|
||||
|
||||
# Student Profile Analyzer Agent
|
||||
profile_analyzer_agent = Agent(
|
||||
agent_name="Student-Profile-Analyzer",
|
||||
system_prompt="""You are an expert student profile analyzer. Your role is to:
|
||||
1. Analyze academic performance, test scores, and extracurricular activities
|
||||
2. Identify student's strengths, weaknesses, and unique qualities
|
||||
3. Evaluate personal statements and essays
|
||||
4. Assess leadership experiences and community involvement
|
||||
5. Determine student's preferences for college environment, location, and programs
|
||||
6. Create a comprehensive student profile summary
|
||||
|
||||
Always consider both quantitative metrics (GPA, test scores) and qualitative aspects
|
||||
(personal growth, challenges overcome, unique perspectives).""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="profile_analyzer_agent.json",
|
||||
user_name="student",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# College Research Agent
|
||||
college_research_agent = Agent(
|
||||
agent_name="College-Research-Specialist",
|
||||
system_prompt="""You are a college research specialist. Your role is to:
|
||||
1. Maintain updated knowledge of college admission requirements
|
||||
2. Research academic programs, campus culture, and student life
|
||||
3. Analyze admission statistics and trends
|
||||
4. Evaluate college-specific opportunities and resources
|
||||
5. Consider financial aid availability and scholarship opportunities
|
||||
6. Track historical admission data and acceptance rates
|
||||
|
||||
Focus on providing accurate, comprehensive information about each institution
|
||||
while considering both academic and cultural fit factors.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="college_research_agent.json",
|
||||
user_name="researcher",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# College Match Agent
|
||||
college_match_agent = Agent(
|
||||
agent_name="College-Match-Maker",
|
||||
system_prompt="""You are a college matching specialist. Your role is to:
|
||||
1. Compare student profiles with college requirements
|
||||
2. Evaluate fit based on academic, social, and cultural factors
|
||||
3. Consider geographic preferences and constraints
|
||||
4. Assess financial fit and aid opportunities
|
||||
5. Create tiered lists of reach, target, and safety schools
|
||||
6. Explain the reasoning behind each match
|
||||
|
||||
Always provide a balanced list with realistic expectations while
|
||||
considering both student preferences and admission probability.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="college_match_agent.json",
|
||||
user_name="matcher",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Debate Moderator Agent
|
||||
debate_moderator_agent = Agent(
|
||||
agent_name="Debate-Moderator",
|
||||
system_prompt="""You are a college selection debate moderator. Your role is to:
|
||||
1. Facilitate discussions between different perspectives
|
||||
2. Ensure all relevant factors are considered
|
||||
3. Challenge assumptions and biases
|
||||
4. Synthesize different viewpoints
|
||||
5. Guide the group toward consensus
|
||||
6. Document key points of agreement and disagreement
|
||||
|
||||
Maintain objectivity while ensuring all important factors are thoroughly discussed
|
||||
and evaluated.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="debate_moderator_agent.json",
|
||||
user_name="moderator",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Critique Agent
|
||||
critique_agent = Agent(
|
||||
agent_name="College-Selection-Critic",
|
||||
system_prompt="""You are a college selection critic. Your role is to:
|
||||
1. Evaluate the strength of college matches
|
||||
2. Identify potential overlooked factors
|
||||
3. Challenge assumptions in the selection process
|
||||
4. Assess risks and potential drawbacks
|
||||
5. Provide constructive feedback on selections
|
||||
6. Suggest alternative options when appropriate
|
||||
|
||||
Focus on constructive criticism that helps improve the final college list
|
||||
while maintaining realistic expectations.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="critique_agent.json",
|
||||
user_name="critic",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Final Decision Agent
|
||||
final_decision_agent = Agent(
|
||||
agent_name="Final-Decision-Maker",
|
||||
system_prompt="""
|
||||
You are a college selection final decision maker. Your role is to:
|
||||
1. Synthesize all previous analyses and discussions
|
||||
2. Weigh competing factors and trade-offs
|
||||
3. Create a final ranked list of recommended colleges
|
||||
4. Provide clear rationale for each recommendation
|
||||
5. Include specific action items for each selected school
|
||||
6. Outline next steps in the application process
|
||||
|
||||
Focus on creating actionable, well-reasoned final recommendations that
|
||||
balance all relevant factors and stakeholder input.
|
||||
""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="final_decision_agent.json",
|
||||
user_name="decision_maker",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Initialize the Sequential Workflow
|
||||
college_selection_workflow = AgentRearrange(
|
||||
name="college-selection-swarm",
|
||||
description="Comprehensive college selection and analysis system",
|
||||
max_loops=1,
|
||||
agents=[
|
||||
profile_analyzer_agent,
|
||||
college_research_agent,
|
||||
college_match_agent,
|
||||
debate_moderator_agent,
|
||||
critique_agent,
|
||||
final_decision_agent,
|
||||
],
|
||||
output_type="all",
|
||||
flow=f"{profile_analyzer_agent.name} -> {college_research_agent.name} -> {college_match_agent.name} -> {debate_moderator_agent.name} -> {critique_agent.name} -> {final_decision_agent.name}",
|
||||
)
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
# Example student profile input
|
||||
student_profile = """
|
||||
Student Profile:
|
||||
- GPA: 3.8
|
||||
- SAT: 1450
|
||||
- Interests: Computer Science, Robotics
|
||||
- Location Preference: East Coast
|
||||
- Extracurriculars: Robotics Club President, Math Team
|
||||
- Budget: Need financial aid
|
||||
- Preferred Environment: Medium-sized urban campus
|
||||
"""
|
||||
|
||||
# Run the comprehensive college selection analysis
|
||||
result = college_selection_workflow.run(
|
||||
student_profile,
|
||||
no_use_clusterops=True,
|
||||
)
|
||||
print(result)
|
||||
@ -0,0 +1,64 @@
|
||||
"""
|
||||
Todo
|
||||
|
||||
- You send structured data to the swarm through the users form they make
|
||||
- then connect rag for every agent using llama index to remember all the students data
|
||||
- structured outputs
|
||||
"""
|
||||
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from swarm_models import OpenAIChat, OpenAIFunctionCaller
|
||||
from pydantic import BaseModel
|
||||
from typing import List
|
||||
|
||||
|
||||
class CollegeLog(BaseModel):
|
||||
college_name: str
|
||||
college_description: str
|
||||
college_admission_requirements: str
|
||||
|
||||
|
||||
class CollegesRecommendation(BaseModel):
|
||||
colleges: List[CollegeLog]
|
||||
reasoning: str
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Get the API key from environment variable
|
||||
api_key = os.getenv("GROQ_API_KEY")
|
||||
|
||||
# Initialize the model
|
||||
model = OpenAIChat(
|
||||
openai_api_base="https://api.groq.com/openai/v1",
|
||||
openai_api_key=api_key,
|
||||
model_name="llama-3.1-70b-versatile",
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
function_caller = OpenAIFunctionCaller(
|
||||
system_prompt="""You are a college selection final decision maker. Your role is to:
|
||||
- Balance all relevant factors and stakeholder input.
|
||||
- Only return the output in the schema format.
|
||||
""",
|
||||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_model=CollegesRecommendation,
|
||||
# parallel_tool_calls=True,
|
||||
)
|
||||
|
||||
|
||||
print(
|
||||
function_caller.run(
|
||||
"""
|
||||
Student Profile: Kye Gomez
|
||||
- GPA: 3.8
|
||||
- SAT: 1450
|
||||
- Interests: Computer Science, Robotics
|
||||
- Location Preference: East Coast
|
||||
- Extracurriculars: Robotics Club President, Math Team
|
||||
- Budget: Need financial aid
|
||||
- Preferred Environment: Medium-sized urban campus
|
||||
"""
|
||||
)
|
||||
)
|
||||
@ -0,0 +1,116 @@
|
||||
from typing import Optional
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
|
||||
|
||||
|
||||
class LlamaIndexDB:
|
||||
"""A class to manage document indexing and querying using LlamaIndex.
|
||||
|
||||
This class provides functionality to add documents from a directory and query the indexed documents.
|
||||
|
||||
Args:
|
||||
data_dir (str): Directory containing documents to index. Defaults to "docs".
|
||||
**kwargs: Additional arguments passed to SimpleDirectoryReader and VectorStoreIndex.
|
||||
SimpleDirectoryReader kwargs:
|
||||
- filename_as_id (bool): Use filenames as document IDs
|
||||
- recursive (bool): Recursively read subdirectories
|
||||
- required_exts (List[str]): Only read files with these extensions
|
||||
- exclude_hidden (bool): Skip hidden files
|
||||
|
||||
VectorStoreIndex kwargs:
|
||||
- service_context: Custom service context
|
||||
- embed_model: Custom embedding model
|
||||
- similarity_top_k (int): Number of similar docs to retrieve
|
||||
- store_nodes_override (bool): Override node storage
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str = "docs", **kwargs) -> None:
|
||||
"""Initialize the LlamaIndexDB with an empty index.
|
||||
|
||||
Args:
|
||||
data_dir (str): Directory containing documents to index
|
||||
**kwargs: Additional arguments for SimpleDirectoryReader and VectorStoreIndex
|
||||
"""
|
||||
self.data_dir = data_dir
|
||||
self.index: Optional[VectorStoreIndex] = None
|
||||
self.reader_kwargs = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k
|
||||
in SimpleDirectoryReader.__init__.__code__.co_varnames
|
||||
}
|
||||
self.index_kwargs = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k not in self.reader_kwargs
|
||||
}
|
||||
|
||||
logger.info("Initialized LlamaIndexDB")
|
||||
data_path = Path(self.data_dir)
|
||||
if not data_path.exists():
|
||||
logger.error(f"Directory not found: {self.data_dir}")
|
||||
raise FileNotFoundError(
|
||||
f"Directory {self.data_dir} does not exist"
|
||||
)
|
||||
|
||||
try:
|
||||
documents = SimpleDirectoryReader(
|
||||
self.data_dir, **self.reader_kwargs
|
||||
).load_data()
|
||||
self.index = VectorStoreIndex.from_documents(
|
||||
documents, **self.index_kwargs
|
||||
)
|
||||
logger.success(
|
||||
f"Successfully indexed documents from {self.data_dir}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error indexing documents: {str(e)}")
|
||||
raise
|
||||
|
||||
def query(self, query: str, **kwargs) -> str:
|
||||
"""Query the indexed documents.
|
||||
|
||||
Args:
|
||||
query (str): The query string to search for
|
||||
**kwargs: Additional arguments passed to the query engine
|
||||
- similarity_top_k (int): Number of similar documents to retrieve
|
||||
- streaming (bool): Enable streaming response
|
||||
- response_mode (str): Response synthesis mode
|
||||
- max_tokens (int): Maximum tokens in response
|
||||
|
||||
Returns:
|
||||
str: The response from the query engine
|
||||
|
||||
Raises:
|
||||
ValueError: If no documents have been indexed yet
|
||||
"""
|
||||
if self.index is None:
|
||||
logger.error("No documents have been indexed yet")
|
||||
raise ValueError("Must add documents before querying")
|
||||
|
||||
try:
|
||||
query_engine = self.index.as_query_engine(**kwargs)
|
||||
response = query_engine.query(query)
|
||||
print(response)
|
||||
logger.info(f"Successfully queried: {query}")
|
||||
return str(response)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during query: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# # Example usage
|
||||
# llama_index_db = LlamaIndexDB(
|
||||
# data_dir="docs",
|
||||
# filename_as_id=True,
|
||||
# recursive=True,
|
||||
# required_exts=[".txt", ".pdf", ".docx"],
|
||||
# similarity_top_k=3
|
||||
# )
|
||||
# response = llama_index_db.query(
|
||||
# "What is the medical history of patient 1?",
|
||||
# streaming=True,
|
||||
# response_mode="compact"
|
||||
# )
|
||||
# print(response)
|
||||
Binary file not shown.
@ -0,0 +1,237 @@
|
||||
"""
|
||||
Todo
|
||||
|
||||
- You send structured data to the swarm through the users form they make
|
||||
- then connect rag for every agent using llama index to remember all the students data
|
||||
- structured outputs
|
||||
"""
|
||||
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from swarms import Agent, SequentialWorkflow
|
||||
from swarm_models import OpenAIChat, OpenAIFunctionCaller
|
||||
from pydantic import BaseModel
|
||||
from typing import List
|
||||
|
||||
|
||||
class CollegeLog(BaseModel):
|
||||
college_name: str
|
||||
college_description: str
|
||||
college_admission_requirements: str
|
||||
|
||||
|
||||
class CollegesRecommendation(BaseModel):
|
||||
colleges: List[CollegeLog]
|
||||
reasoning: str
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Get the API key from environment variable
|
||||
api_key = os.getenv("GROQ_API_KEY")
|
||||
|
||||
# Initialize the model
|
||||
model = OpenAIChat(
|
||||
openai_api_base="https://api.groq.com/openai/v1",
|
||||
openai_api_key=api_key,
|
||||
model_name="llama-3.1-70b-versatile",
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
FINAL_AGENT_PROMPT = """
|
||||
You are a college selection final decision maker. Your role is to:
|
||||
1. Synthesize all previous analyses and discussions
|
||||
2. Weigh competing factors and trade-offs
|
||||
3. Create a final ranked list of recommended colleges
|
||||
4. Provide clear rationale for each recommendation
|
||||
5. Include specific action items for each selected school
|
||||
6. Outline next steps in the application process
|
||||
|
||||
Focus on creating actionable, well-reasoned final recommendations that
|
||||
balance all relevant factors and stakeholder input.
|
||||
|
||||
"""
|
||||
|
||||
function_caller = OpenAIFunctionCaller(
|
||||
system_prompt=FINAL_AGENT_PROMPT,
|
||||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_model=CollegesRecommendation,
|
||||
parallel_tool_calls=True,
|
||||
)
|
||||
|
||||
# Student Profile Analyzer Agent
|
||||
profile_analyzer_agent = Agent(
|
||||
agent_name="Student-Profile-Analyzer",
|
||||
system_prompt="""You are an expert student profile analyzer. Your role is to:
|
||||
1. Analyze academic performance, test scores, and extracurricular activities
|
||||
2. Identify student's strengths, weaknesses, and unique qualities
|
||||
3. Evaluate personal statements and essays
|
||||
4. Assess leadership experiences and community involvement
|
||||
5. Determine student's preferences for college environment, location, and programs
|
||||
6. Create a comprehensive student profile summary
|
||||
|
||||
Always consider both quantitative metrics (GPA, test scores) and qualitative aspects
|
||||
(personal growth, challenges overcome, unique perspectives).""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="profile_analyzer_agent.json",
|
||||
user_name="student",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# College Research Agent
|
||||
college_research_agent = Agent(
|
||||
agent_name="College-Research-Specialist",
|
||||
system_prompt="""You are a college research specialist. Your role is to:
|
||||
1. Maintain updated knowledge of college admission requirements
|
||||
2. Research academic programs, campus culture, and student life
|
||||
3. Analyze admission statistics and trends
|
||||
4. Evaluate college-specific opportunities and resources
|
||||
5. Consider financial aid availability and scholarship opportunities
|
||||
6. Track historical admission data and acceptance rates
|
||||
|
||||
Focus on providing accurate, comprehensive information about each institution
|
||||
while considering both academic and cultural fit factors.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="college_research_agent.json",
|
||||
user_name="researcher",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# College Match Agent
|
||||
college_match_agent = Agent(
|
||||
agent_name="College-Match-Maker",
|
||||
system_prompt="""You are a college matching specialist. Your role is to:
|
||||
1. Compare student profiles with college requirements
|
||||
2. Evaluate fit based on academic, social, and cultural factors
|
||||
3. Consider geographic preferences and constraints
|
||||
4. Assess financial fit and aid opportunities
|
||||
5. Create tiered lists of reach, target, and safety schools
|
||||
6. Explain the reasoning behind each match
|
||||
|
||||
Always provide a balanced list with realistic expectations while
|
||||
considering both student preferences and admission probability.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="college_match_agent.json",
|
||||
user_name="matcher",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Debate Moderator Agent
|
||||
debate_moderator_agent = Agent(
|
||||
agent_name="Debate-Moderator",
|
||||
system_prompt="""You are a college selection debate moderator. Your role is to:
|
||||
1. Facilitate discussions between different perspectives
|
||||
2. Ensure all relevant factors are considered
|
||||
3. Challenge assumptions and biases
|
||||
4. Synthesize different viewpoints
|
||||
5. Guide the group toward consensus
|
||||
6. Document key points of agreement and disagreement
|
||||
|
||||
Maintain objectivity while ensuring all important factors are thoroughly discussed
|
||||
and evaluated.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="debate_moderator_agent.json",
|
||||
user_name="moderator",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Critique Agent
|
||||
critique_agent = Agent(
|
||||
agent_name="College-Selection-Critic",
|
||||
system_prompt="""You are a college selection critic. Your role is to:
|
||||
1. Evaluate the strength of college matches
|
||||
2. Identify potential overlooked factors
|
||||
3. Challenge assumptions in the selection process
|
||||
4. Assess risks and potential drawbacks
|
||||
5. Provide constructive feedback on selections
|
||||
6. Suggest alternative options when appropriate
|
||||
|
||||
Focus on constructive criticism that helps improve the final college list
|
||||
while maintaining realistic expectations.""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="critique_agent.json",
|
||||
user_name="critic",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Final Decision Agent
|
||||
final_decision_agent = Agent(
|
||||
agent_name="Final-Decision-Maker",
|
||||
system_prompt="""
|
||||
You are a college selection final decision maker. Your role is to:
|
||||
1. Synthesize all previous analyses and discussions
|
||||
2. Weigh competing factors and trade-offs
|
||||
3. Create a final ranked list of recommended colleges
|
||||
4. Provide clear rationale for each recommendation
|
||||
5. Include specific action items for each selected school
|
||||
6. Outline next steps in the application process
|
||||
|
||||
Focus on creating actionable, well-reasoned final recommendations that
|
||||
balance all relevant factors and stakeholder input.
|
||||
""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="final_decision_agent.json",
|
||||
user_name="decision_maker",
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Initialize the Sequential Workflow
|
||||
college_selection_workflow = SequentialWorkflow(
|
||||
name="college-selection-swarm",
|
||||
description="Comprehensive college selection and analysis system",
|
||||
max_loops=1,
|
||||
agents=[
|
||||
profile_analyzer_agent,
|
||||
college_research_agent,
|
||||
college_match_agent,
|
||||
debate_moderator_agent,
|
||||
critique_agent,
|
||||
final_decision_agent,
|
||||
],
|
||||
output_type="all",
|
||||
)
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
# Example student profile input
|
||||
student_profile = """
|
||||
Student Profile:
|
||||
- GPA: 3.8
|
||||
- SAT: 1450
|
||||
- Interests: Computer Science, Robotics
|
||||
- Location Preference: East Coast
|
||||
- Extracurriculars: Robotics Club President, Math Team
|
||||
- Budget: Need financial aid
|
||||
- Preferred Environment: Medium-sized urban campus
|
||||
"""
|
||||
|
||||
# Run the comprehensive college selection analysis
|
||||
result = college_selection_workflow.run(
|
||||
student_profile,
|
||||
no_use_clusterops=True,
|
||||
)
|
||||
print(result)
|
||||
@ -0,0 +1,143 @@
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from swarms import Agent, SequentialWorkflow
|
||||
from swarm_models import OpenAIChat
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Get the OpenAI API key from the environment variable
|
||||
api_key = os.getenv("GROQ_API_KEY")
|
||||
|
||||
# Model
|
||||
model = OpenAIChat(
|
||||
openai_api_base="https://api.groq.com/openai/v1",
|
||||
openai_api_key=api_key,
|
||||
model_name="llama-3.1-70b-versatile",
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
|
||||
# Initialize specialized agents
|
||||
data_extractor_agent = Agent(
|
||||
agent_name="Data-Extractor",
|
||||
system_prompt="""You are a data extraction specialist. Your role is to:
|
||||
1. Extract key information, data points, and metrics from documents
|
||||
2. Identify and pull out important facts, figures, and statistics
|
||||
3. Structure extracted data in a clear, organized format
|
||||
4. Flag any inconsistencies or missing data
|
||||
5. Ensure accuracy in data extraction while maintaining context""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="data_extractor_agent.json",
|
||||
user_name="pe_firm",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
summarizer_agent = Agent(
|
||||
agent_name="Document-Summarizer",
|
||||
system_prompt="""You are a document summarization expert. Your role is to:
|
||||
1. Create concise, comprehensive summaries of documents
|
||||
2. Highlight key points and main takeaways
|
||||
3. Maintain the essential meaning while reducing length
|
||||
4. Structure summaries in a logical, readable format
|
||||
5. Identify and emphasize critical insights""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="summarizer_agent.json",
|
||||
user_name="pe_firm",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
financial_analyst_agent = Agent(
|
||||
agent_name="Financial-Analyst",
|
||||
system_prompt="""You are a financial analysis expert. Your role is to:
|
||||
1. Analyze financial statements and metrics
|
||||
2. Evaluate company valuations and financial projections
|
||||
3. Assess financial risks and opportunities
|
||||
4. Provide insights on financial performance and health
|
||||
5. Make recommendations based on financial analysis""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="financial_analyst_agent.json",
|
||||
user_name="pe_firm",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
market_analyst_agent = Agent(
|
||||
agent_name="Market-Analyst",
|
||||
system_prompt="""You are a market analysis expert. Your role is to:
|
||||
1. Analyze market trends and dynamics
|
||||
2. Evaluate competitive landscape and market positioning
|
||||
3. Identify market opportunities and threats
|
||||
4. Assess market size and growth potential
|
||||
5. Provide strategic market insights and recommendations""",
|
||||
llm=model,
|
||||
max_loops=1,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="market_analyst_agent.json",
|
||||
user_name="pe_firm",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
operational_analyst_agent = Agent(
|
||||
agent_name="Operational-Analyst",
|
||||
system_prompt="""You are an operational analysis expert. Your role is to:
|
||||
1. Analyze business operations and processes
|
||||
2. Evaluate operational efficiency and effectiveness
|
||||
3. Identify operational risks and opportunities
|
||||
4. Assess scalability and growth potential
|
||||
5. Provide recommendations for operational improvements""",
|
||||
llm=model,
|
||||
max_loops=2,
|
||||
autosave=True,
|
||||
verbose=True,
|
||||
dynamic_temperature_enabled=True,
|
||||
saved_state_path="operational_analyst_agent.json",
|
||||
user_name="pe_firm",
|
||||
retry_attempts=1,
|
||||
context_length=200000,
|
||||
output_type="string",
|
||||
)
|
||||
|
||||
# Initialize the SwarmRouter
|
||||
router = SequentialWorkflow(
|
||||
name="pe-document-analysis-swarm",
|
||||
description="Analyze documents for private equity due diligence and investment decision-making",
|
||||
max_loops=1,
|
||||
agents=[
|
||||
data_extractor_agent,
|
||||
summarizer_agent,
|
||||
financial_analyst_agent,
|
||||
market_analyst_agent,
|
||||
operational_analyst_agent,
|
||||
],
|
||||
output_type="all",
|
||||
)
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
# Run a comprehensive private equity document analysis task
|
||||
result = router.run(
|
||||
"Where is the best place to find template term sheets for series A startups. Provide links and references",
|
||||
no_use_clusterops=True,
|
||||
)
|
||||
print(result)
|
||||
@ -0,0 +1,135 @@
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List
|
||||
|
||||
from rich.console import Console
|
||||
from rich.live import Live
|
||||
from rich.panel import Panel
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn
|
||||
from rich.table import Table
|
||||
from rich.text import Text
|
||||
|
||||
|
||||
class Formatter:
|
||||
"""
|
||||
A class for formatting and printing rich text to the console.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Initializes the Formatter with a Rich Console instance.
|
||||
"""
|
||||
self.console = Console()
|
||||
|
||||
def print_panel(
|
||||
self, content: str, title: str = "", style: str = "bold blue"
|
||||
) -> None:
|
||||
"""
|
||||
Prints a rich panel to the console with a random color.
|
||||
|
||||
Args:
|
||||
content (str): The content of the panel.
|
||||
title (str, optional): The title of the panel. Defaults to "".
|
||||
style (str, optional): The style of the panel. Defaults to "bold blue".
|
||||
"""
|
||||
import random
|
||||
|
||||
colors = [
|
||||
"red",
|
||||
"green",
|
||||
"blue",
|
||||
"yellow",
|
||||
"magenta",
|
||||
"cyan",
|
||||
"white",
|
||||
]
|
||||
random_color = random.choice(colors)
|
||||
panel = Panel(
|
||||
content, title=title, style=f"bold {random_color}"
|
||||
)
|
||||
self.console.print(panel)
|
||||
|
||||
def print_table(
|
||||
self, title: str, data: Dict[str, List[str]]
|
||||
) -> None:
|
||||
"""
|
||||
Prints a rich table to the console.
|
||||
|
||||
Args:
|
||||
title (str): The title of the table.
|
||||
data (Dict[str, List[str]]): A dictionary where keys are categories and values are lists of capabilities.
|
||||
"""
|
||||
table = Table(show_header=True, header_style="bold magenta")
|
||||
table.add_column("Category", style="cyan")
|
||||
table.add_column("Capabilities", style="green")
|
||||
|
||||
for category, items in data.items():
|
||||
table.add_row(category, "\n".join(items))
|
||||
|
||||
self.console.print(f"\n🔥 {title}:", style="bold yellow")
|
||||
self.console.print(table)
|
||||
|
||||
def print_progress(
|
||||
self,
|
||||
description: str,
|
||||
task_fn: Callable,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""
|
||||
Prints a progress bar to the console and executes a task function.
|
||||
|
||||
Args:
|
||||
description (str): The description of the task.
|
||||
task_fn (Callable): The function to execute.
|
||||
*args (Any): Arguments to pass to the task function.
|
||||
**kwargs (Any): Keyword arguments to pass to the task function.
|
||||
|
||||
Returns:
|
||||
Any: The result of the task function.
|
||||
"""
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
) as progress:
|
||||
task = progress.add_task(description, total=None)
|
||||
result = task_fn(*args, **kwargs)
|
||||
progress.update(task, completed=True)
|
||||
return result
|
||||
|
||||
def print_panel_token_by_token(
|
||||
self,
|
||||
tokens: str,
|
||||
title: str = "Output",
|
||||
style: str = "bold cyan",
|
||||
delay: float = 0.01,
|
||||
by_word: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Prints a string in real-time, token by token (character or word) inside a Rich panel.
|
||||
|
||||
Args:
|
||||
tokens (str): The string to display in real-time.
|
||||
title (str): Title of the panel.
|
||||
style (str): Style for the text inside the panel.
|
||||
delay (float): Delay in seconds between displaying each token.
|
||||
by_word (bool): If True, display by words; otherwise, display by characters.
|
||||
"""
|
||||
text = Text(style=style)
|
||||
|
||||
# Split tokens into characters or words
|
||||
token_list = tokens.split() if by_word else tokens
|
||||
|
||||
with Live(
|
||||
Panel(text, title=title, border_style=style),
|
||||
console=self.console,
|
||||
refresh_per_second=10,
|
||||
) as live:
|
||||
for token in token_list:
|
||||
text.append(token + (" " if by_word else ""))
|
||||
live.update(
|
||||
Panel(text, title=title, border_style=style)
|
||||
)
|
||||
time.sleep(delay)
|
||||
|
||||
|
||||
formatter = Formatter()
|
||||
Loading…
Reference in new issue