from siqse import LLMSuggestionSelector

# Initialize the user simulator
simulator = LLMSuggestionSelector(
    model="llama3.1:70b",
    llm_api_url="http://localhost:11434",
    selection_mode="batch",
    use_feedback_docs=False
)

# Simulate which suggestions a user would select
query = "machine learning applications"
suggestions = ["neural networks", "deep learning", "AI ethics", "data mining"]

selected = simulator.select_suggestions(
    query=query,
    suggestions=suggestions,
    n=2  # Simulate selecting top 2
)

print(selected)  # e.g., ["neural networks", "deep learning"]

def get_retrieval_docs(query: str, selected_keywords: list) -> list:
    """
    Retrieve documents to show to the simulated user.
    Returns list of dicts with keys: doc_id, score, rank, is_relevant, text
    """
    # Your retrieval logic here
    return [
        {
            "doc_id": "doc1",
            "score": 0.95,
            "rank": 1,
            "is_relevant": True,
            "text": "Document text here..."
        },
        # ... more documents
    ]

# Simulate iterative user behavior with feedback
simulator = LLMSuggestionSelector(
    model="llama3.1:70b",
    selection_mode="iterative",
    use_feedback_docs=True
)

selected = simulator.select_suggestions(
    query=query,
    suggestions=suggestions,
    user_intent="Find recent research papers on neural network architectures",
    n=3,
    get_docs=get_retrieval_docs
)

# Simulate user deciding when to stop selecting
selected = simulator.select_suggestions(
    query=query,
    suggestions=suggestions,
    n=None  # No limit - simulated user decides when to stop
)

simulator = LLMSuggestionSelector(selection_mode="batch")

simulator = LLMSuggestionSelector(selection_mode="iterative")

import pyterrier as pt
from siqse import LLMSuggestionSelector

# Initialize
if not pt.started():
    pt.init()

index = pt.IndexFactory.of("/path/to/index")
simulator = LLMSuggestionSelector(
    model="llama3.1:8b",
    selection_mode="iterative",
    use_feedback_docs=True
)
bm25 = pt.BatchRetrieve(index, wmodel="BM25")

# Define expansion function that simulates user selection
def expand(df):
    def fb(q, s):
        # Show results to simulated user for decision-making
        expanded = q + " " + " ".join(s) if s else q
        return bm25.search(expanded).head(10).to_dict('records')
    
    df = df.copy()
    df["query"] = df.apply(
        lambda r: " ".join([r["query"]] + simulator.select_suggestions(
            query=r["query"],
            suggestions=r["suggestions"],
            user_intent=r.get("intent"),
            get_docs=fb
        )),
        axis=1
    )
    return df

# Build pipeline: simulate user selection then retrieve
pipeline = pt.apply.generic(expand) >> bm25
results = pipeline.transform(queries_df)

from siqse import (
    clean_query_for_pyterrier,
    format_document_for_display,
    count_relevant_documents
)

# Clean query for IR systems
clean_query = clean_query_for_pyterrier("Parkinson's disease")

# Format document for LLM prompt
doc = {
    "doc_id": "doc1",
    "score": 0.95,
    "rank": 1,
    "is_relevant": True,
    "text": "Long document text..."
}
formatted = format_document_for_display(doc, snippet_chars=300)

# Count relevant docs
docs = [...]  # list of doc dicts
num_relevant = count_relevant_documents(docs)