expand + knowledge graph

knowledge.py

@@ -31,12 +31,13 @@ class KnowledgeGraph(BaseModel):
 
 from groq import Groq
 import os
-os.environ["PATH"] += os.pathsep + '/usr/bin/'
+
 # Initialize with API key
 client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 
 # Enable instructor patches for Groq client
 client = instructor.from_groq(client)
+llm='llama-3.1-8b-instant' #"llama3.2", #
 """
 from openai import OpenAI
 client = instructor.from_openai(
@@ -47,18 +48,46 @@ client = instructor.from_openai(
     mode=instructor.Mode.JSON,
 )
 """
-def generate_graph(input) -> KnowledgeGraph:
+def generate_graph(q, input) -> KnowledgeGraph:
     return client.chat.completions.create(
-    model='llama-3.1-8b-instant', #"llama3.2", #
+    model=llm,
     max_retries=5,
     messages=[
         {
             "role": "user",
-               "content": f"Help me understand the following by describing it as a detailed knowledge graph: {input}",
+            "content": f"Help me understand the following by describing it as a detailed knowledge graph: ### Question: {q} ### Context: {input}",
         }
     ],
     response_model=KnowledgeGraph,
 )
+
+
+class Issue(BaseModel):
+    "Break down Issue as sub issues"
+    question: str
+
+class IssuePlan(BaseModel):
+    "List of Issue"
+    issue_graph: List[Issue] = []
+
+
+def expandIssue(input) -> IssuePlan:
+    return client.chat.completions.create(
+        model=llm,
+        max_retries=10,
+        messages=[
+            {
+                "role": "system",
+                "content": "As a Mckinsey Consultant create a framework that relevant to the topic, list all issues.",
+            },             {
+                "role": "user",
+                "content": f"Question: {input}",
+            },
+        ],
+        response_model=IssuePlan,
+    )
 def graph(query):
-    graph = generate_graph(query)
+    queryx = expandIssue(query)
+    ctx = ", ".join([q.question for q in queryx.issue_graph])
+    graph = generate_graph(query, ctx)
     return graph.json()