from modaic import PrecompiledAgent, PrecompiledConfig

class WeatherConfig(PrecompiledConfig):
    weather: str = "sunny"

class WeatherAgent(PrecompiledAgent):
    config: WeatherConfig

    def __init__(self, config: WeatherConfig, **kwargs):
        super().__init__(config, **kwargs)

    def forward(self, query: str) -> str:
        return f"The weather in {query} is {self.config.weather}."

agent = WeatherAgent(WeatherConfig())
print(agent(query="Tokyo"))

agent.save_precompiled("./my-weather")

from modaic import AutoAgent, AutoConfig

cfg = AutoConfig.from_precompiled("./my-weather", local=True)
loaded = AutoAgent.from_precompiled("./my-weather", local=True)
print(loaded(query="Kyoto"))

from pathlib import Path
from modaic.context import Table, TableFile
import pandas as pd

# Load from Excel/CSV
excel = TableFile.from_file(
    file_ref="employees.xlsx",
    file=Path("employees.xlsx"),
    file_type="xlsx",
)
csv = TableFile.from_file(
    file_ref="data.csv",
    file=Path("data.csv"),
    file_type="csv",
)

# Create from DataFrame
df = pd.DataFrame({"col1": [1, 2, 3], "col2": [4, 5, 6]})
table = Table(df=df, name="my_table")

# Query with SQL (refer to in-memory table as `this`)
result = table.query("SELECT * FROM this WHERE col1 > 1")

# Convert to markdown
markdown = table.markdown()

from modaic.databases import SQLDatabase, SQLiteBackend

# Configure and connect
backend = SQLiteBackend(db_path="my_database.db")
db = SQLDatabase(backend)

# Add table
db.add_table(table)

# Query
rows = db.fetchall("SELECT * FROM my_table")

from modaic.context import Context, Relation
from modaic.databases import GraphDatabase, MemgraphConfig, Neo4jConfig

# Configure backend (choose one)
mg = GraphDatabase(MemgraphConfig())
# or
neo = GraphDatabase(Neo4jConfig())

# Define nodes
class Person(Context):
    name: str
    age: int

class KNOWS(Relation):
    since: int

alice = Person(name="Alice", age=30)
bob = Person(name="Bob", age=28)

# Save nodes
alice.save(mg)
bob.save(mg)

# Create relationship (Alice)-[KNOWS]->(Bob)
rel = (alice >> KNOWS(since=2020) >> bob)
rel.save(mg)

# Query
rows = mg.execute_and_fetch("MATCH (a:Person)-[r:KNOWS]->(b:Person) RETURN a, r, b LIMIT 5")

from modaic import Embedder
from modaic.context import Text
from modaic.databases import VectorDatabase, MilvusBackend

# Setup embedder and backend
embedder = Embedder("openai/text-embedding-3-small")
backend = MilvusBackend.from_local("vector.db")  # milvus lite

# Initialize database
vdb = VectorDatabase(backend=backend, embedder=embedder, payload_class=Text)

# Create collection and add records
vdb.create_collection("my_collection", payload_class=Text)
vdb.add_records("my_collection", [Text(text="hello world"), Text(text="modaic makes sharing agents easy")])

# Search
results = vdb.search("my_collection", query="hello", k=3)
top_hit_text = results[0][0].context.text

from modaic.precompiled_agent import PrecompiledConfig, PrecompiledAgent
from modaic.context import Table
from modaic.databases import VectorDatabase, SQLDatabase
from modaic.types import Indexer

class TableRAGConfig(PrecompiledConfig):
    k_recall: int = 50
    k_rerank: int = 5

class TableRAGAgent(PrecompiledAgent):
    config: TableRAGConfig # ! Important: config must be annotated with the config class
    
    def __init__(self, config: TableRAGConfig, indexer: Indexer, **kwargs):
        super().__init__(config, **kwargs)
        self.indexer = indexer
        # Initialize DSPy modules for reasoning
    
    def forward(self, user_query: str) -> str:
        # Retrieve relevant tables
        # Generate SQL queries
        # Combine results and provide answer
        pass