File size: 4,665 Bytes
b81ac13
 
860db05
b81ac13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
860db05
 
 
 
 
 
b81ac13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import VectorStoreIndex, Document
# from llama_index.llms.openllm import OpenLLM
from llama_index.llms.nebius import NebiusLLM
import requests
import os

# Load environment variables from .env if present
from dotenv import load_dotenv
load_dotenv()

# Read provider, keys, and model names from environment
LLM_PROVIDER = os.environ.get("LLM_PROVIDER", "openllm").lower()
LLM_API_URL = os.environ.get("LLM_API_URL")
LLM_API_KEY = os.environ.get("LLM_API_KEY")
NEBIUS_API_KEY = os.environ.get("NEBIUS_API_KEY", "")
OPENLLM_MODEL = os.environ.get("OPENLLM_MODEL", "neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16")
NEBIUS_MODEL = os.environ.get("NEBIUS_MODEL", "meta-llama/Llama-3.3-70B-Instruct")

# Choose LLM provider
if LLM_PROVIDER == "nebius":
    llm = NebiusLLM(
        api_key=NEBIUS_API_KEY,
        model=NEBIUS_MODEL
    )
else:
    pass
    # llm = OpenLLM(
    #     model=OPENLLM_MODEL,
    #     api_base=LLM_API_URL,
    #     api_key=LLM_API_KEY
    # )   

# Example: Define your stages and their descriptions here
STAGE_DOCS = [
    Document(text="Goal setting: Define what you want to achieve."),
    Document(text="Research: Gather information and resources."),
    Document(text="Planning: Break down your goal into actionable steps."),
    Document(text="Execution: Start working on your plan."),
    Document(text="Review: Reflect on your progress and adjust as needed."),
]

# Stage-specific instructions for each stage
STAGE_INSTRUCTIONS = {
    "Goal setting": (
        "After trying to understand the goal, before moving to the next phase, "
        "write down key objectives that the user is interested in."
    ),
    "Research": (
        "Before suggesting something to the user, think deeply about what scientific approach you are using to suggest something or ask a question. "
        "Before moving to a new phase, summarize in a detailed format the key findings of research and intuition."
    ),
    "Planning": (
        "Provide a detailed actionable plan with a proper timeline. "
        "Try to create tasks in 3 types: Important and have a deadline, Important but do not have a timeline, Not important and has a deadline."
    ),
    "Execution": (
        "Focus on helping the user execute the plan step by step. Offer encouragement and practical advice."
    ),
    "Review": (
        "Help the user reflect on progress, identify what worked, and suggest adjustments for future improvement."
    ),
}

def get_stage_instruction(stage_name):
    """
    Returns the instruction string for a given stage name, or an empty string if not found.
    """
    return STAGE_INSTRUCTIONS.get(stage_name, "")

def build_index():
    embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
    # Always build the index from the documents, so text is present
    return VectorStoreIndex.from_documents(STAGE_DOCS, embed_model=embed_model)

# Build the index once (reuse for all queries)
index = build_index()

def map_stage(user_input):
    # Use your custom LLM for generative responses if needed
    query_engine = index.as_query_engine(similarity_top_k=1, llm=llm)
    response = query_engine.query(user_input)
    # Return the most relevant stage and its details
    return {
        "stage": response.source_nodes[0].node.text,
        "details": response.response
    }

def get_stage_and_details(user_input):
    """
    Helper to get stage and details for a given user input.
    """
    query_engine = index.as_query_engine(similarity_top_k=1, llm=llm)
    response = query_engine.query(user_input)
    stage = response.source_nodes[0].node.text
    details = response.response
    return stage, details

def clear_vector_store():
    if os.path.exists(VECTOR_STORE_PATH):
        os.remove(VECTOR_STORE_PATH)

def get_stage_list():
    """
    Returns the ordered list of stage names.
    """
    return [
        "Goal setting",
        "Research",
        "Planning",
        "Execution",
        "Review"
    ]

def get_next_stage(current_stage):
    """
    Given the current stage name, returns the next stage name or None if at the end.
    """
    stages = get_stage_list()
    try:
        idx = stages.index(current_stage)
        if idx + 1 < len(stages):
            return stages[idx + 1]
    except ValueError:
        pass
    return None

def get_stage_index(stage_name):
    """
    Returns the index of the given stage name in the ordered list, or -1 if not found.
    """
    try:
        return get_stage_list().index(stage_name)
    except ValueError:
        return -1